/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.hadoop.hbase.client;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;

import java.io.IOException;
import java.io.InterruptedIOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.CallDroppedException;
import org.apache.hadoop.hbase.CallQueueTooBigException;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.HBaseClassTestRule;
import org.apache.hadoop.hbase.HBaseIOException;
import org.apache.hadoop.hbase.HBaseServerException;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.HRegionInfo;
import org.apache.hadoop.hbase.HRegionLocation;
import org.apache.hadoop.hbase.RegionLocations;
import org.apache.hadoop.hbase.ServerName;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.AsyncProcessTask.ListRowAccess;
import org.apache.hadoop.hbase.client.AsyncProcessTask.SubmittedRows;
import org.apache.hadoop.hbase.client.backoff.ClientBackoffPolicy;
import org.apache.hadoop.hbase.client.backoff.ServerStatistics;
import org.apache.hadoop.hbase.client.coprocessor.Batch;
import org.apache.hadoop.hbase.exceptions.RegionOpeningException;
import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
import org.apache.hadoop.hbase.quotas.RpcThrottlingException;
import org.apache.hadoop.hbase.security.User;
import org.apache.hadoop.hbase.testclassification.ClientTests;
import org.apache.hadoop.hbase.testclassification.LargeTests;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
import org.apache.hadoop.hbase.util.Threads;
import org.junit.Assert;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.hbase.thirdparty.com.google.common.util.concurrent.ThreadFactoryBuilder;

@Category({ ClientTests.class, LargeTests.class })
public class TestAsyncProcess {

  @ClassRule
  public static final HBaseClassTestRule CLASS_RULE =
    HBaseClassTestRule.forClass(TestAsyncProcess.class);

  private static final Logger LOG = LoggerFactory.getLogger(TestAsyncProcess.class);
  private static final TableName DUMMY_TABLE = TableName.valueOf("DUMMY_TABLE");
  private static final byte[] DUMMY_BYTES_1 = Bytes.toBytes("DUMMY_BYTES_1");
  private static final byte[] DUMMY_BYTES_2 = Bytes.toBytes("DUMMY_BYTES_2");
  private static final byte[] DUMMY_BYTES_3 = Bytes.toBytes("DUMMY_BYTES_3");
  private static final byte[] FAILS = Bytes.toBytes("FAILS");
  private Configuration CONF;
  private ConnectionConfiguration CONNECTION_CONFIG;
  private static final ServerName sn = ServerName.valueOf("s1,1,1");
  private static final ServerName sn2 = ServerName.valueOf("s2,2,2");
  private static final ServerName sn3 = ServerName.valueOf("s3,3,3");
  private static final HRegionInfo hri1 =
    new HRegionInfo(DUMMY_TABLE, DUMMY_BYTES_1, DUMMY_BYTES_2, false, 1);
  private static final HRegionInfo hri2 =
    new HRegionInfo(DUMMY_TABLE, DUMMY_BYTES_2, HConstants.EMPTY_END_ROW, false, 2);
  private static final HRegionInfo hri3 =
    new HRegionInfo(DUMMY_TABLE, DUMMY_BYTES_3, HConstants.EMPTY_END_ROW, false, 3);
  private static final HRegionLocation loc1 = new HRegionLocation(hri1, sn);
  private static final HRegionLocation loc2 = new HRegionLocation(hri2, sn);
  private static final HRegionLocation loc3 = new HRegionLocation(hri3, sn2);

  // Replica stuff
  private static final RegionInfo hri1r1 = RegionReplicaUtil.getRegionInfoForReplica(hri1, 1);
  private static final RegionInfo hri1r2 = RegionReplicaUtil.getRegionInfoForReplica(hri1, 2);
  private static final RegionInfo hri2r1 = RegionReplicaUtil.getRegionInfoForReplica(hri2, 1);
  private static final RegionLocations hrls1 = new RegionLocations(new HRegionLocation(hri1, sn),
    new HRegionLocation(hri1r1, sn2), new HRegionLocation(hri1r2, sn3));
  private static final RegionLocations hrls2 =
    new RegionLocations(new HRegionLocation(hri2, sn2), new HRegionLocation(hri2r1, sn3));
  private static final RegionLocations hrls3 =
    new RegionLocations(new HRegionLocation(hri3, sn3), null);

  private static final String success = "success";
  private static Exception failure = new Exception("failure");

  private static final int NB_RETRIES = 3;

  private int RPC_TIMEOUT;
  private int OPERATION_TIMEOUT;

  @Before
  public void beforeEach() {
    this.CONF = new Configuration();
    CONF.setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, NB_RETRIES);
    this.CONNECTION_CONFIG = new ConnectionConfiguration(CONF);
    this.RPC_TIMEOUT =
      CONF.getInt(HConstants.HBASE_RPC_TIMEOUT_KEY, HConstants.DEFAULT_HBASE_RPC_TIMEOUT);
    this.OPERATION_TIMEOUT = CONF.getInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT,
      HConstants.DEFAULT_HBASE_CLIENT_OPERATION_TIMEOUT);
  }

  static class CountingThreadFactory implements ThreadFactory {
    final AtomicInteger nbThreads;
    ThreadFactory realFactory =
      new ThreadFactoryBuilder().setNameFormat("test-TestAsyncProcess-pool-%d")
        .setUncaughtExceptionHandler(Threads.LOGGING_EXCEPTION_HANDLER).build();

    @Override
    public Thread newThread(Runnable r) {
      nbThreads.incrementAndGet();
      return realFactory.newThread(r);
    }

    CountingThreadFactory(AtomicInteger nbThreads) {
      this.nbThreads = nbThreads;
    }
  }

  static class MyAsyncProcess extends AsyncProcess {
    final AtomicInteger nbMultiResponse = new AtomicInteger();
    final AtomicInteger nbActions = new AtomicInteger();
    public List<AsyncRequestFuture> allReqs = new ArrayList<>();
    public AtomicInteger callsCt = new AtomicInteger();
    private Configuration conf;

    private long previousTimeout = -1;
    final ExecutorService service;

    @Override
    protected <Res> AsyncRequestFutureImpl<Res> createAsyncRequestFuture(AsyncProcessTask task,
      List<Action> actions, long nonceGroup) {
      // Test HTable has tableName of null, so pass DUMMY_TABLE
      AsyncProcessTask wrap = new AsyncProcessTask(task) {
        @Override
        public TableName getTableName() {
          return DUMMY_TABLE;
        }
      };
      AsyncRequestFutureImpl<Res> r =
        new MyAsyncRequestFutureImpl<>(wrap, actions, nonceGroup, this);
      allReqs.add(r);
      return r;
    }

    public MyAsyncProcess(ClusterConnection hc, Configuration conf) {
      super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
        new RpcControllerFactory(conf));
      service = Executors.newFixedThreadPool(5);
      this.conf = conf;
    }

    public MyAsyncProcess(ClusterConnection hc, Configuration conf, AtomicInteger nbThreads) {
      super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
        new RpcControllerFactory(conf));
      service = new ThreadPoolExecutor(1, 20, 60, TimeUnit.SECONDS, new SynchronousQueue<>(),
        new CountingThreadFactory(nbThreads));
    }

    public <CResult> AsyncRequestFuture submit(ExecutorService pool, TableName tableName,
      List<? extends Row> rows, boolean atLeastOne, Batch.Callback<CResult> callback,
      boolean needResults) throws InterruptedIOException {
      AsyncProcessTask task = AsyncProcessTask.newBuilder(callback)
        .setPool(pool == null ? service : pool).setTableName(tableName).setRowAccess(rows)
        .setSubmittedRows(atLeastOne ? SubmittedRows.AT_LEAST_ONE : SubmittedRows.NORMAL)
        .setNeedResults(needResults)
        .setRpcTimeout(
          conf.getInt(HConstants.HBASE_RPC_READ_TIMEOUT_KEY, HConstants.DEFAULT_HBASE_RPC_TIMEOUT))
        .setOperationTimeout(conf.getInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT,
          HConstants.DEFAULT_HBASE_CLIENT_OPERATION_TIMEOUT))
        .build();
      return submit(task);
    }

    public <CResult> AsyncRequestFuture submit(TableName tableName, final List<? extends Row> rows,
      boolean atLeastOne, Batch.Callback<CResult> callback, boolean needResults)
      throws InterruptedIOException {
      return submit(null, tableName, rows, atLeastOne, callback, needResults);
    }

    @Override
    public <Res> AsyncRequestFuture submit(AsyncProcessTask<Res> task)
      throws InterruptedIOException {
      previousTimeout = task.getRpcTimeout();
      // We use results in tests to check things, so override to always save them.
      AsyncProcessTask<Res> wrap = new AsyncProcessTask<Res>(task) {
        @Override
        public boolean getNeedResults() {
          return true;
        }
      };
      return super.submit(wrap);
    }

    @Override
    protected RpcRetryingCaller<AbstractResponse>
      createCaller(CancellableRegionServerCallable callable, int rpcTimeout) {
      callsCt.incrementAndGet();
      MultiServerCallable callable1 = (MultiServerCallable) callable;
      final MultiResponse mr = createMultiResponse(callable1.getMulti(), nbMultiResponse, nbActions,
        new ResponseGenerator() {
          @Override
          public void addResponse(MultiResponse mr, byte[] regionName, Action a) {
            if (Arrays.equals(FAILS, a.getAction().getRow())) {
              mr.add(regionName, a.getOriginalIndex(), failure);
            } else {
              mr.add(regionName, a.getOriginalIndex(), success);
            }
          }
        });

      return new RpcRetryingCallerImpl<AbstractResponse>(100, 500, 10,
        RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, 9, 0, null) {
        @Override
        public AbstractResponse callWithoutRetries(RetryingCallable<AbstractResponse> callable,
          int callTimeout) throws IOException, RuntimeException {
          try {
            // sleep one second in order for threadpool to start another thread instead of reusing
            // existing one.
            Thread.sleep(1000);
          } catch (InterruptedException e) {
            // ignore error
          }
          return mr;
        }
      };
    }

  }

  static class MyAsyncRequestFutureImpl<Res> extends AsyncRequestFutureImpl<Res> {
    private final Map<ServerName, List<Long>> heapSizesByServer = new HashMap<>();

    public MyAsyncRequestFutureImpl(AsyncProcessTask task, List<Action> actions, long nonceGroup,
      AsyncProcess asyncProcess) {
      super(task, actions, nonceGroup, asyncProcess);
    }

    @Override
    protected void updateStats(ServerName server, MultiResponse resp) {
      // Do nothing for avoiding the NPE if we test the ClientBackofPolicy.
    }

    Map<ServerName, List<Long>> getRequestHeapSize() {
      return heapSizesByServer;
    }

    @Override
    SingleServerRequestRunnable createSingleServerRequest(MultiAction multiAction, int numAttempt,
      ServerName server, Set<CancellableRegionServerCallable> callsInProgress) {
      SingleServerRequestRunnable rq =
        new SingleServerRequestRunnable(multiAction, numAttempt, server, callsInProgress);
      List<Long> heapCount = heapSizesByServer.get(server);
      if (heapCount == null) {
        heapCount = new ArrayList<>();
        heapSizesByServer.put(server, heapCount);
      }
      heapCount.add(heapSizeOf(multiAction));
      return rq;
    }

    private long heapSizeOf(MultiAction multiAction) {
      return multiAction.actions.values().stream().flatMap(v -> v.stream())
        .map(action -> action.getAction()).filter(row -> row instanceof Mutation)
        .mapToLong(row -> ((Mutation) row).heapSize()).sum();
    }
  }

  static class CallerWithFailure extends RpcRetryingCallerImpl<AbstractResponse> {

    private final IOException e;

    public CallerWithFailure(IOException e) {
      super(100, 500, 100, RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, 9, 0, null);
      this.e = e;
    }

    @Override
    public AbstractResponse callWithoutRetries(RetryingCallable<AbstractResponse> callable,
      int callTimeout) throws IOException, RuntimeException {
      throw e;
    }
  }

  /**
   * Used to simulate the case where a RegionServer responds to a multi request, but some or all of
   * the actions have an Exception instead of Result. These responses go through receiveMultiAction,
   * which has handling for individual action failures.
   */
  static class CallerWithRegionException extends RpcRetryingCallerImpl<AbstractResponse> {

    private final IOException e;
    private MultiAction multi;

    public CallerWithRegionException(IOException e, MultiAction multi) {
      super(100, 500, 100, RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, 9, 0, null);
      this.e = e;
      this.multi = multi;
    }

    @Override
    public AbstractResponse callWithoutRetries(RetryingCallable<AbstractResponse> callable,
      int callTimeout) throws IOException, RuntimeException {
      MultiResponse response = new MultiResponse();
      for (Entry<byte[], List<Action>> entry : multi.actions.entrySet()) {
        response.addException(entry.getKey(), e);
      }
      return response;
    }
  }

  static class AsyncProcessWithFailure extends MyAsyncProcess {

    private final IOException ioe;
    private final ServerName failingServer;
    private final boolean returnAsRegionException;

    public AsyncProcessWithFailure(ClusterConnection hc, Configuration myConf, IOException ioe,
      ServerName failingServer, boolean returnAsRegionException) {
      super(hc, myConf);
      this.ioe = ioe;
      this.failingServer = failingServer;
      this.returnAsRegionException = returnAsRegionException;
      serverTrackerTimeout = 1L;
    }

    public AsyncProcessWithFailure(ClusterConnection hc, Configuration myConf, IOException ioe) {
      this(hc, myConf, ioe, null, false);
    }

    @Override
    protected RpcRetryingCaller<AbstractResponse>
      createCaller(CancellableRegionServerCallable callable, int rpcTimeout) {
      MultiServerCallable msc = (MultiServerCallable) callable;
      if (failingServer != null) {
        if (!msc.getServerName().equals(failingServer)) {
          return super.createCaller(callable, rpcTimeout);
        }
      }

      if (returnAsRegionException) {
        return new CallerWithRegionException(ioe, msc.getMulti());
      }

      callsCt.incrementAndGet();
      return new CallerWithFailure(ioe);
    }
  }

  /**
   * Make the backoff time always different on each call.
   */
  static class MyClientBackoffPolicy implements ClientBackoffPolicy {
    private final Map<ServerName, AtomicInteger> count = new HashMap<>();

    @Override
    public long getBackoffTime(ServerName serverName, byte[] region, ServerStatistics stats) {
      AtomicInteger inc = count.get(serverName);
      if (inc == null) {
        inc = new AtomicInteger(0);
        count.put(serverName, inc);
      }
      return inc.getAndIncrement();
    }
  }

  static class MyAsyncProcessWithReplicas extends MyAsyncProcess {
    private Set<byte[]> failures = new TreeSet<>(new Bytes.ByteArrayComparator());
    private long primarySleepMs = 0, replicaSleepMs = 0;
    private Map<ServerName, Long> customPrimarySleepMs = new HashMap<>();
    private final AtomicLong replicaCalls = new AtomicLong(0);

    public void addFailures(RegionInfo... hris) {
      for (RegionInfo hri : hris) {
        failures.add(hri.getRegionName());
      }
    }

    public long getReplicaCallCount() {
      return replicaCalls.get();
    }

    public void setPrimaryCallDelay(ServerName server, long primaryMs) {
      customPrimarySleepMs.put(server, primaryMs);
    }

    public MyAsyncProcessWithReplicas(ClusterConnection hc, Configuration conf) {
      super(hc, conf);
    }

    public void setCallDelays(long primaryMs, long replicaMs) {
      this.primarySleepMs = primaryMs;
      this.replicaSleepMs = replicaMs;
    }

    @Override
    protected RpcRetryingCaller<AbstractResponse>
      createCaller(CancellableRegionServerCallable payloadCallable, int rpcTimeout) {
      MultiServerCallable callable = (MultiServerCallable) payloadCallable;
      final MultiResponse mr = createMultiResponse(callable.getMulti(), nbMultiResponse, nbActions,
        new ResponseGenerator() {
          @Override
          public void addResponse(MultiResponse mr, byte[] regionName, Action a) {
            if (failures.contains(regionName)) {
              mr.add(regionName, a.getOriginalIndex(), failure);
            } else {
              boolean isStale = !RegionReplicaUtil.isDefaultReplica(a.getReplicaId());
              mr.add(regionName, a.getOriginalIndex(), Result.create(new Cell[0], null, isStale));
            }
          }
        });
      // Currently AsyncProcess either sends all-replica, or all-primary request.
      final boolean isDefault = RegionReplicaUtil.isDefaultReplica(
        callable.getMulti().actions.values().iterator().next().iterator().next().getReplicaId());
      final ServerName server = ((MultiServerCallable) callable).getServerName();
      String debugMsg = "Call to " + server + ", primary=" + isDefault + " with "
        + callable.getMulti().actions.size() + " entries: ";
      for (byte[] region : callable.getMulti().actions.keySet()) {
        debugMsg += "[" + Bytes.toStringBinary(region) + "], ";
      }
      LOG.debug(debugMsg);
      if (!isDefault) {
        replicaCalls.incrementAndGet();
      }

      return new RpcRetryingCallerImpl<AbstractResponse>(100, 500, 10,
        RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, 9, 0, null) {
        @Override
        public MultiResponse callWithoutRetries(RetryingCallable<AbstractResponse> callable,
          int callTimeout) throws IOException, RuntimeException {
          long sleep = -1;
          if (isDefault) {
            Long customSleep = customPrimarySleepMs.get(server);
            sleep = (customSleep == null ? primarySleepMs : customSleep.longValue());
          } else {
            sleep = replicaSleepMs;
          }
          if (sleep != 0) {
            try {
              Thread.sleep(sleep);
            } catch (InterruptedException e) {
              // Restore interrupt status
              Thread.currentThread().interrupt();
            }
          }
          return mr;
        }
      };
    }
  }

  static MultiResponse createMultiResponse(final MultiAction multi, AtomicInteger nbMultiResponse,
    AtomicInteger nbActions, ResponseGenerator gen) {
    final MultiResponse mr = new MultiResponse();
    nbMultiResponse.incrementAndGet();
    for (Map.Entry<byte[], List<Action>> entry : multi.actions.entrySet()) {
      byte[] regionName = entry.getKey();
      for (Action a : entry.getValue()) {
        nbActions.incrementAndGet();
        gen.addResponse(mr, regionName, a);
      }
    }
    return mr;
  }

  private static interface ResponseGenerator {
    void addResponse(final MultiResponse mr, byte[] regionName, Action a);
  }

  /**
   * Returns our async process.
   */
  static class MyConnectionImpl extends ConnectionImplementation {
    public static class TestRegistry extends DoNothingConnectionRegistry {

      public TestRegistry(Configuration conf, User user) {
        super(conf, user);
      }

      @Override
      public CompletableFuture<String> getClusterId() {
        return CompletableFuture.completedFuture("testClusterId");
      }
    }

    final AtomicInteger nbThreads = new AtomicInteger(0);

    protected MyConnectionImpl(Configuration conf) throws IOException {
      super(setupConf(conf), null, null, Collections.emptyMap());
    }

    private static Configuration setupConf(Configuration conf) {
      conf.setClass(HConstants.CLIENT_CONNECTION_REGISTRY_IMPL_CONF_KEY, TestRegistry.class,
        ConnectionRegistry.class);
      return conf;
    }

    @Override
    public RegionLocations locateRegion(TableName tableName, byte[] row, boolean useCache,
      boolean retry, int replicaId) throws IOException {
      return new RegionLocations(loc1);
    }

    @Override
    public boolean hasCellBlockSupport() {
      return false;
    }
  }

  /**
   * Returns our async process.
   */
  static class MyConnectionImpl2 extends MyConnectionImpl {
    List<HRegionLocation> hrl;
    final boolean usedRegions[];

    protected MyConnectionImpl2(List<HRegionLocation> hrl, Configuration conf) throws IOException {
      super(conf);
      this.hrl = hrl;
      this.usedRegions = new boolean[hrl.size()];
    }

    @Override
    public RegionLocations locateRegion(TableName tableName, byte[] row, boolean useCache,
      boolean retry, int replicaId) throws IOException {
      int i = 0;
      for (HRegionLocation hr : hrl) {
        if (Arrays.equals(row, hr.getRegionInfo().getStartKey())) {
          usedRegions[i] = true;
          return new RegionLocations(hr);
        }
        i++;
      }
      return null;
    }
  }

  @Test
  public void testListRowAccess() {
    int count = 10;
    List<String> values = new ArrayList<>();
    for (int i = 0; i != count; ++i) {
      values.add(String.valueOf(i));
    }

    ListRowAccess<String> taker = new ListRowAccess<>(values);
    assertEquals(count, taker.size());

    int takeCount = 0;
    Iterator<String> it = taker.iterator();
    while (it.hasNext()) {
      String v = it.next();
      assertEquals(String.valueOf(takeCount), v);
      ++takeCount;
      it.remove();
      if (Math.random() >= 0.5) {
        break;
      }
    }
    assertEquals(count, taker.size() + takeCount);

    it = taker.iterator();
    while (it.hasNext()) {
      String v = it.next();
      assertEquals(String.valueOf(takeCount), v);
      ++takeCount;
      it.remove();
    }
    assertEquals(0, taker.size());
    assertEquals(count, takeCount);
  }

  private static long calculateRequestCount(long putSizePerServer, long maxHeapSizePerRequest) {
    if (putSizePerServer <= maxHeapSizePerRequest) {
      return 1;
    } else if (putSizePerServer % maxHeapSizePerRequest == 0) {
      return putSizePerServer / maxHeapSizePerRequest;
    } else {
      return putSizePerServer / maxHeapSizePerRequest + 1;
    }
  }

  @Test
  public void testSubmitSameSizeOfRequest() throws Exception {
    long writeBuffer = 2 * 1024 * 1024;
    long putsHeapSize = writeBuffer;
    doSubmitRequest(writeBuffer, putsHeapSize);
  }

  @Test
  public void testSubmitLargeRequestWithUnlimitedSize() throws Exception {
    long maxHeapSizePerRequest = Long.MAX_VALUE;
    long putsHeapSize = 2 * 1024 * 1024;
    doSubmitRequest(maxHeapSizePerRequest, putsHeapSize);
  }

  @Test
  public void testSubmitRandomSizeRequest() throws Exception {
    Random rn = new Random();
    final long limit = 10 * 1024 * 1024;
    final int requestCount = 1 + (int) (rn.nextDouble() * 3);
    long n = rn.nextLong();
    if (n < 0) {
      n = -n;
    } else if (n == 0) {
      n = 1;
    }
    long putsHeapSize = n % limit;
    long maxHeapSizePerRequest = putsHeapSize / requestCount;
    LOG.info("[testSubmitRandomSizeRequest] maxHeapSizePerRequest=" + maxHeapSizePerRequest
      + ", putsHeapSize=" + putsHeapSize);
    doSubmitRequest(maxHeapSizePerRequest, putsHeapSize);
  }

  @Test
  public void testSubmitSmallRequest() throws Exception {
    long maxHeapSizePerRequest = 2 * 1024 * 1024;
    long putsHeapSize = 100;
    doSubmitRequest(maxHeapSizePerRequest, putsHeapSize);
  }

  @Test
  public void testSubmitLargeRequest() throws Exception {
    long maxHeapSizePerRequest = 2 * 1024 * 1024;
    long putsHeapSize = maxHeapSizePerRequest * 2;
    doSubmitRequest(maxHeapSizePerRequest, putsHeapSize);
  }

  private void doSubmitRequest(long maxHeapSizePerRequest, long putsHeapSize) throws Exception {
    ClusterConnection conn = createHConnection();
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    final long defaultHeapSizePerRequest =
      conn.getConfiguration().getLong(SimpleRequestController.HBASE_CLIENT_MAX_PERREQUEST_HEAPSIZE,
        SimpleRequestController.DEFAULT_HBASE_CLIENT_MAX_PERREQUEST_HEAPSIZE);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    conn.getConfiguration().setLong(SimpleRequestController.HBASE_CLIENT_MAX_PERREQUEST_HEAPSIZE,
      maxHeapSizePerRequest);

    // sn has two regions
    long putSizeSN = 0;
    long putSizeSN2 = 0;
    List<Put> puts = new ArrayList<>();
    while ((putSizeSN + putSizeSN2) <= putsHeapSize) {
      Put put1 = new Put(DUMMY_BYTES_1);
      put1.addColumn(DUMMY_BYTES_1, DUMMY_BYTES_1, DUMMY_BYTES_1);
      Put put2 = new Put(DUMMY_BYTES_2);
      put2.addColumn(DUMMY_BYTES_2, DUMMY_BYTES_2, DUMMY_BYTES_2);
      Put put3 = new Put(DUMMY_BYTES_3);
      put3.addColumn(DUMMY_BYTES_3, DUMMY_BYTES_3, DUMMY_BYTES_3);
      putSizeSN += (put1.heapSize() + put2.heapSize());
      putSizeSN2 += put3.heapSize();
      puts.add(put1);
      puts.add(put2);
      puts.add(put3);
    }

    int minCountSnRequest = (int) calculateRequestCount(putSizeSN, maxHeapSizePerRequest);
    int minCountSn2Request = (int) calculateRequestCount(putSizeSN2, maxHeapSizePerRequest);
    LOG.info("Total put count:" + puts.size() + ", putSizeSN:" + putSizeSN + ", putSizeSN2:"
      + putSizeSN2 + ", maxHeapSizePerRequest:" + maxHeapSizePerRequest + ", minCountSnRequest:"
      + minCountSnRequest + ", minCountSn2Request:" + minCountSn2Request);

    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    try (BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap)) {
      mutator.mutate(puts);
      mutator.flush();
      List<AsyncRequestFuture> reqs = ap.allReqs;

      int actualSnReqCount = 0;
      int actualSn2ReqCount = 0;
      for (AsyncRequestFuture req : reqs) {
        if (!(req instanceof AsyncRequestFutureImpl)) {
          continue;
        }
        MyAsyncRequestFutureImpl ars = (MyAsyncRequestFutureImpl) req;
        if (ars.getRequestHeapSize().containsKey(sn)) {
          ++actualSnReqCount;
        }
        if (ars.getRequestHeapSize().containsKey(sn2)) {
          ++actualSn2ReqCount;
        }
      }
      // If the server is busy, the actual count may be incremented.
      assertEquals(true, minCountSnRequest <= actualSnReqCount);
      assertEquals(true, minCountSn2Request <= actualSn2ReqCount);
      Map<ServerName, Long> sizePerServers = new HashMap<>();
      for (AsyncRequestFuture req : reqs) {
        if (!(req instanceof AsyncRequestFutureImpl)) {
          continue;
        }
        MyAsyncRequestFutureImpl ars = (MyAsyncRequestFutureImpl) req;
        Map<ServerName, List<Long>> requestHeapSize = ars.getRequestHeapSize();
        for (Map.Entry<ServerName, List<Long>> entry : requestHeapSize.entrySet()) {
          long sum = 0;
          for (long size : entry.getValue()) {
            assertEquals(true, size <= maxHeapSizePerRequest);
            sum += size;
          }
          assertEquals(true, sum <= maxHeapSizePerRequest);
          long value = sizePerServers.getOrDefault(entry.getKey(), 0L);
          sizePerServers.put(entry.getKey(), value + sum);
        }
      }
      assertEquals(true, sizePerServers.containsKey(sn));
      assertEquals(true, sizePerServers.containsKey(sn2));
      assertEquals(false, sizePerServers.containsKey(sn3));
      assertEquals(putSizeSN, (long) sizePerServers.get(sn));
      assertEquals(putSizeSN2, (long) sizePerServers.get(sn2));
    }
    // restore config.
    conn.getConfiguration().setLong(SimpleRequestController.HBASE_CLIENT_MAX_PERREQUEST_HEAPSIZE,
      defaultHeapSizePerRequest);
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  @Test
  public void testSubmit() throws Exception {
    ClusterConnection hc = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(hc, CONF);

    List<Put> puts = new ArrayList<>(1);
    puts.add(createPut(1, true));

    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertTrue(puts.isEmpty());
  }

  @Test
  public void testSubmitWithCB() throws Exception {
    ClusterConnection hc = createHConnection();
    final AtomicInteger updateCalled = new AtomicInteger(0);
    Batch.Callback<Object> cb = new Batch.Callback<Object>() {
      @Override
      public void update(byte[] region, byte[] row, Object result) {
        updateCalled.incrementAndGet();
      }
    };
    MyAsyncProcess ap = new MyAsyncProcess(hc, CONF);

    List<Put> puts = new ArrayList<>(1);
    puts.add(createPut(1, true));

    final AsyncRequestFuture ars = ap.submit(null, DUMMY_TABLE, puts, false, cb, false);
    Assert.assertTrue(puts.isEmpty());
    ars.waitUntilDone();
    Assert.assertEquals(1, updateCalled.get());
  }

  @Test
  public void testSubmitBusyRegion() throws Exception {
    ClusterConnection conn = createHConnection();
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    SimpleRequestController controller = (SimpleRequestController) ap.requestController;
    List<Put> puts = new ArrayList<>(1);
    puts.add(createPut(1, true));

    for (int i = 0; i != controller.maxConcurrentTasksPerRegion; ++i) {
      ap.incTaskCounters(Collections.singleton(hri1.getRegionName()), sn);
    }
    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertEquals(puts.size(), 1);

    ap.decTaskCounters(Collections.singleton(hri1.getRegionName()), sn);
    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertEquals(0, puts.size());
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  @Test
  public void testSubmitBusyRegionServer() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    SimpleRequestController controller = (SimpleRequestController) ap.requestController;
    controller.taskCounterPerServer.put(sn2,
      new AtomicInteger(controller.maxConcurrentTasksPerServer));

    List<Put> puts = new ArrayList<>(4);
    puts.add(createPut(1, true));
    puts.add(createPut(3, true)); // <== this one won't be taken, the rs is busy
    puts.add(createPut(1, true)); // <== this one will make it, the region is already in
    puts.add(createPut(2, true)); // <== new region, but the rs is ok

    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertEquals(" puts=" + puts, 1, puts.size());

    controller.taskCounterPerServer.put(sn2,
      new AtomicInteger(controller.maxConcurrentTasksPerServer - 1));
    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertTrue(puts.isEmpty());
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  @Test
  public void testFail() throws Exception {
    MyAsyncProcess ap = new MyAsyncProcess(createHConnection(), CONF);

    List<Put> puts = new ArrayList<>(1);
    Put p = createPut(1, false);
    puts.add(p);

    AsyncRequestFuture ars = ap.submit(null, DUMMY_TABLE, puts, false, null, true);
    Assert.assertEquals(0, puts.size());
    ars.waitUntilDone();
    verifyResult(ars, false);
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());

    Assert.assertEquals(1, ars.getErrors().exceptions.size());
    Assert.assertTrue("was: " + ars.getErrors().exceptions.get(0),
      failure.equals(ars.getErrors().exceptions.get(0)));
    Assert.assertTrue("was: " + ars.getErrors().exceptions.get(0),
      failure.equals(ars.getErrors().exceptions.get(0)));

    Assert.assertEquals(1, ars.getFailedOperations().size());
    Assert.assertTrue("was: " + ars.getFailedOperations().get(0),
      p.equals(ars.getFailedOperations().get(0)));
  }

  @Test
  public void testSubmitTrue() throws IOException {
    ClusterConnection conn = createHConnection();
    final MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    SimpleRequestController controller = (SimpleRequestController) ap.requestController;
    controller.tasksInProgress.incrementAndGet();
    final AtomicInteger ai = new AtomicInteger(controller.maxConcurrentTasksPerRegion);
    controller.taskCounterPerRegion.put(hri1.getRegionName(), ai);

    final AtomicBoolean checkPoint = new AtomicBoolean(false);
    final AtomicBoolean checkPoint2 = new AtomicBoolean(false);

    Thread t = new Thread() {
      @Override
      public void run() {
        Threads.sleep(1000);
        Assert.assertFalse(checkPoint.get()); // TODO: this is timing-dependent
        ai.decrementAndGet();
        controller.tasksInProgress.decrementAndGet();
        checkPoint2.set(true);
      }
    };

    List<Put> puts = new ArrayList<>(1);
    Put p = createPut(1, true);
    puts.add(p);

    ap.submit(null, DUMMY_TABLE, puts, false, null, false);
    Assert.assertFalse(puts.isEmpty());

    t.start();

    ap.submit(null, DUMMY_TABLE, puts, true, null, false);
    Assert.assertTrue(puts.isEmpty());

    checkPoint.set(true);
    while (!checkPoint2.get()) {
      Threads.sleep(1);
    }
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  @Test
  public void testFailAndSuccess() throws Exception {
    MyAsyncProcess ap = new MyAsyncProcess(createHConnection(), CONF);

    List<Put> puts = new ArrayList<>(3);
    puts.add(createPut(1, false));
    puts.add(createPut(1, true));
    puts.add(createPut(1, true));

    AsyncRequestFuture ars = ap.submit(null, DUMMY_TABLE, puts, false, null, true);
    Assert.assertTrue(puts.isEmpty());
    ars.waitUntilDone();
    verifyResult(ars, false, true, true);
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());
    ap.callsCt.set(0);
    Assert.assertEquals(1, ars.getErrors().actions.size());

    puts.add(createPut(1, true));
    // Wait for AP to be free. While ars might have the result, ap counters are decreased later.
    ap.waitForMaximumCurrentTasks(0, null);
    ars = ap.submit(null, DUMMY_TABLE, puts, false, null, true);
    Assert.assertEquals(0, puts.size());
    ars.waitUntilDone();
    Assert.assertEquals(1, ap.callsCt.get());
    verifyResult(ars, true);
  }

  @Test
  public void testFlush() throws Exception {
    MyAsyncProcess ap = new MyAsyncProcess(createHConnection(), CONF);

    List<Put> puts = new ArrayList<>(3);
    puts.add(createPut(1, false));
    puts.add(createPut(1, true));
    puts.add(createPut(1, true));

    AsyncRequestFuture ars = ap.submit(null, DUMMY_TABLE, puts, false, null, true);
    ars.waitUntilDone();
    verifyResult(ars, false, true, true);
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());

    Assert.assertEquals(1, ars.getFailedOperations().size());
  }

  @Test
  public void testTaskCountWithoutClientBackoffPolicy() throws IOException, InterruptedException {
    ClusterConnection hc = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(hc, CONF);
    testTaskCount(ap);
  }

  @Test
  public void testTaskCountWithClientBackoffPolicy() throws IOException, InterruptedException {
    Configuration copyConf = new Configuration(CONF);
    copyConf.setBoolean(HConstants.ENABLE_CLIENT_BACKPRESSURE, true);
    MyClientBackoffPolicy bp = new MyClientBackoffPolicy();
    ClusterConnection conn = createHConnection();
    Mockito.when(conn.getConfiguration()).thenReturn(copyConf);
    Mockito.when(conn.getStatisticsTracker()).thenReturn(ServerStatisticTracker.create(copyConf));
    Mockito.when(conn.getBackoffPolicy()).thenReturn(bp);
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    MyAsyncProcess ap = new MyAsyncProcess(conn, copyConf);
    testTaskCount(ap);
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  private void testTaskCount(MyAsyncProcess ap)
    throws InterruptedIOException, InterruptedException {
    SimpleRequestController controller = (SimpleRequestController) ap.requestController;
    List<Put> puts = new ArrayList<>();
    for (int i = 0; i != 3; ++i) {
      puts.add(createPut(1, true));
      puts.add(createPut(2, true));
      puts.add(createPut(3, true));
    }
    ap.submit(null, DUMMY_TABLE, puts, true, null, false);
    ap.waitForMaximumCurrentTasks(0, null);
    // More time to wait if there are incorrect task count.
    TimeUnit.SECONDS.sleep(1);
    assertEquals(0, controller.tasksInProgress.get());
    for (AtomicInteger count : controller.taskCounterPerRegion.values()) {
      assertEquals(0, count.get());
    }
    for (AtomicInteger count : controller.taskCounterPerServer.values()) {
      assertEquals(0, count.get());
    }
  }

  @Test
  public void testMaxTask() throws Exception {
    ClusterConnection conn = createHConnection();
    final String defaultClazz =
      conn.getConfiguration().get(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY);
    conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
      SimpleRequestController.class.getName());
    final MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    SimpleRequestController controller = (SimpleRequestController) ap.requestController;

    for (int i = 0; i < 1000; i++) {
      ap.incTaskCounters(Collections.singleton(Bytes.toBytes("dummy")), sn);
    }

    final Thread myThread = Thread.currentThread();

    Thread t = new Thread() {
      @Override
      public void run() {
        Threads.sleep(2000);
        myThread.interrupt();
      }
    };

    List<Put> puts = new ArrayList<>(1);
    puts.add(createPut(1, true));

    t.start();

    try {
      ap.submit(null, DUMMY_TABLE, puts, false, null, false);
      Assert.fail("We should have been interrupted.");
    } catch (InterruptedIOException expected) {
    }

    final long sleepTime = 2000;

    Thread t2 = new Thread() {
      @Override
      public void run() {
        Threads.sleep(sleepTime);
        while (controller.tasksInProgress.get() > 0) {
          ap.decTaskCounters(Collections.singleton(Bytes.toBytes("dummy")), sn);
        }
      }
    };
    t2.start();

    long start = EnvironmentEdgeManager.currentTime();
    ap.submit(null, DUMMY_TABLE, new ArrayList<>(), false, null, false);
    long end = EnvironmentEdgeManager.currentTime();

    // Adds 100 to secure us against approximate timing.
    Assert.assertTrue(start + 100L + sleepTime > end);
    if (defaultClazz != null) {
      conn.getConfiguration().set(RequestControllerFactory.REQUEST_CONTROLLER_IMPL_CONF_KEY,
        defaultClazz);
    }
  }

  private ClusterConnection createHConnection() throws IOException {
    return createHConnection(CONNECTION_CONFIG);
  }

  private ClusterConnection createHConnection(ConnectionConfiguration configuration)
    throws IOException {
    ClusterConnection hc = createHConnectionCommon(configuration);
    setMockLocation(hc, DUMMY_BYTES_1, new RegionLocations(loc1));
    setMockLocation(hc, DUMMY_BYTES_2, new RegionLocations(loc2));
    setMockLocation(hc, DUMMY_BYTES_3, new RegionLocations(loc3));
    Mockito
      .when(hc.locateRegions(Mockito.eq(DUMMY_TABLE), Mockito.anyBoolean(), Mockito.anyBoolean()))
      .thenReturn(Arrays.asList(loc1, loc2, loc3));
    setMockLocation(hc, FAILS, new RegionLocations(loc2));
    return hc;
  }

  private ClusterConnection createHConnectionWithReplicas(ConnectionConfiguration configuration)
    throws IOException {
    ClusterConnection hc = createHConnectionCommon(configuration);
    setMockLocation(hc, DUMMY_BYTES_1, hrls1);
    setMockLocation(hc, DUMMY_BYTES_2, hrls2);
    setMockLocation(hc, DUMMY_BYTES_3, hrls3);
    List<HRegionLocation> locations = new ArrayList<>();
    for (HRegionLocation loc : hrls1.getRegionLocations()) {
      locations.add(loc);
    }
    for (HRegionLocation loc : hrls2.getRegionLocations()) {
      locations.add(loc);
    }
    for (HRegionLocation loc : hrls3.getRegionLocations()) {
      locations.add(loc);
    }
    Mockito
      .when(hc.locateRegions(Mockito.eq(DUMMY_TABLE), Mockito.anyBoolean(), Mockito.anyBoolean()))
      .thenReturn(locations);
    return hc;
  }

  private static void setMockLocation(ClusterConnection hc, byte[] row, RegionLocations result)
    throws IOException {
    Mockito.when(hc.locateRegion(Mockito.eq(DUMMY_TABLE), Mockito.eq(row), Mockito.anyBoolean(),
      Mockito.anyBoolean(), Mockito.anyInt())).thenReturn(result);
    Mockito.when(hc.locateRegion(Mockito.eq(DUMMY_TABLE), Mockito.eq(row), Mockito.anyBoolean(),
      Mockito.anyBoolean())).thenReturn(result);
  }

  private ClusterConnection
    createHConnectionCommon(ConnectionConfiguration connectionConfiguration) {
    ClusterConnection hc = Mockito.mock(ClusterConnection.class);
    NonceGenerator ng = Mockito.mock(NonceGenerator.class);
    Mockito.when(ng.getNonceGroup()).thenReturn(HConstants.NO_NONCE);
    Mockito.when(hc.getNonceGenerator()).thenReturn(ng);
    Mockito.when(hc.getConfiguration()).thenReturn(CONF);
    Mockito.when(hc.getConnectionConfiguration()).thenReturn(connectionConfiguration);
    return hc;
  }

  @Test
  public void testHTablePutSuccess() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl ht = new BufferedMutatorImpl(conn, bufferParam, ap);

    Put put = createPut(1, true);

    Assert.assertEquals(conn.getConnectionConfiguration().getWriteBufferSize(),
      ht.getWriteBufferSize());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());
    ht.mutate(put);
    ht.flush();
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());
  }

  @Test
  public void testSettingWriteBufferPeriodicFlushParameters() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);

    checkPeriodicFlushParameters(conn, ap, 1234, 1234, 1234, 1234);
    checkPeriodicFlushParameters(conn, ap, 0, 0, 0,
      BufferedMutator.MIN_WRITE_BUFFER_PERIODIC_FLUSH_TIMERTICK_MS);
    checkPeriodicFlushParameters(conn, ap, -1234, 0, -1234,
      BufferedMutator.MIN_WRITE_BUFFER_PERIODIC_FLUSH_TIMERTICK_MS);
    checkPeriodicFlushParameters(conn, ap, 1, 1, 1,
      BufferedMutator.MIN_WRITE_BUFFER_PERIODIC_FLUSH_TIMERTICK_MS);
  }

  private void checkPeriodicFlushParameters(ClusterConnection conn, MyAsyncProcess ap, long setTO,
    long expectTO, long setTT, long expectTT) {
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);

    // The BufferedMutatorParams does nothing with the value
    bufferParam.setWriteBufferPeriodicFlushTimeoutMs(setTO);
    bufferParam.setWriteBufferPeriodicFlushTimerTickMs(setTT);
    Assert.assertEquals(setTO, bufferParam.getWriteBufferPeriodicFlushTimeoutMs());
    Assert.assertEquals(setTT, bufferParam.getWriteBufferPeriodicFlushTimerTickMs());

    // The BufferedMutatorImpl corrects illegal values (indirect via BufferedMutatorParams)
    BufferedMutatorImpl ht1 = new BufferedMutatorImpl(conn, bufferParam, ap);
    Assert.assertEquals(expectTO, ht1.getWriteBufferPeriodicFlushTimeoutMs());
    Assert.assertEquals(expectTT, ht1.getWriteBufferPeriodicFlushTimerTickMs());

    // The BufferedMutatorImpl corrects illegal values (direct via setter)
    BufferedMutatorImpl ht2 =
      new BufferedMutatorImpl(conn, createBufferedMutatorParams(ap, DUMMY_TABLE), ap);
    ht2.setWriteBufferPeriodicFlush(setTO, setTT);
    Assert.assertEquals(expectTO, ht2.getWriteBufferPeriodicFlushTimeoutMs());
    Assert.assertEquals(expectTT, ht2.getWriteBufferPeriodicFlushTimerTickMs());

  }

  @Test
  public void testWriteBufferPeriodicFlushTimeoutMs() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);

    bufferParam.setWriteBufferPeriodicFlushTimeoutMs(1); // Flush ASAP
    bufferParam.setWriteBufferPeriodicFlushTimerTickMs(1); // Check every 100ms
    bufferParam.writeBufferSize(10000); // Write buffer set to much larger than the single record

    BufferedMutatorImpl ht = new BufferedMutatorImpl(conn, bufferParam, ap);

    // Verify if BufferedMutator has the right settings.
    Assert.assertEquals(10000, ht.getWriteBufferSize());
    Assert.assertEquals(1, ht.getWriteBufferPeriodicFlushTimeoutMs());
    Assert.assertEquals(BufferedMutator.MIN_WRITE_BUFFER_PERIODIC_FLUSH_TIMERTICK_MS,
      ht.getWriteBufferPeriodicFlushTimerTickMs());

    Put put = createPut(1, true);

    Assert.assertEquals(0, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());

    // ----- Insert, flush immediately, MUST NOT flush automatically
    ht.mutate(put);
    ht.flush();

    Thread.sleep(1000);
    Assert.assertEquals(0, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());

    // ----- Insert, NO flush, MUST flush automatically
    ht.mutate(put);
    Assert.assertEquals(0, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertTrue(ht.getCurrentWriteBufferSize() > 0);

    // The timerTick should fire every 100ms, so after twice that we must have
    // seen at least 1 tick and we should see an automatic flush
    Thread.sleep(200);
    Assert.assertEquals(1, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());

    // Ensure it does not flush twice
    Thread.sleep(200);
    Assert.assertEquals(1, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());

    // ----- DISABLE AUTO FLUSH, Insert, NO flush, MUST NOT flush automatically
    ht.disableWriteBufferPeriodicFlush();
    ht.mutate(put);
    Assert.assertEquals(1, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertTrue(ht.getCurrentWriteBufferSize() > 0);

    // Wait for at least 1 timerTick, we should see NO flushes.
    Thread.sleep(200);
    Assert.assertEquals(1, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertTrue(ht.getCurrentWriteBufferSize() > 0);

    // Reenable periodic flushing, a flush seems to take about 1 second
    // so we wait for 2 seconds and it should have finished the flush.
    ht.setWriteBufferPeriodicFlush(1, 100);
    Thread.sleep(2000);
    Assert.assertEquals(2, ht.getExecutedWriteBufferPeriodicFlushes());
    Assert.assertEquals(0, ht.getCurrentWriteBufferSize());
  }

  @Test
  public void testBufferedMutatorImplWithSharedPool() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutator ht = new BufferedMutatorImpl(conn, bufferParam, ap);

    ht.close();
    assertFalse(ap.service.isShutdown());
  }

  @Test
  public void testFailedPutAndNewPut() throws Exception {
    ClusterConnection conn = createHConnection();
    MyAsyncProcess ap = new MyAsyncProcess(conn, CONF);
    BufferedMutatorParams bufferParam =
      createBufferedMutatorParams(ap, DUMMY_TABLE).writeBufferSize(0);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);

    Put p = createPut(1, false);
    try {
      mutator.mutate(p);
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    // Let's do all the retries.
    ap.waitForMaximumCurrentTasks(0, null);
    Assert.assertEquals(0, mutator.size());

    // There is no global error so the new put should not fail
    mutator.mutate(createPut(1, true));
    Assert.assertEquals("the put should not been inserted.", 0, mutator.size());
  }

  @SuppressWarnings("SelfComparison")
  @Test
  public void testAction() {
    Action action_0 = new Action(new Put(Bytes.toBytes("abc")), 10);
    Action action_1 = new Action(new Put(Bytes.toBytes("ccc")), 10);
    Action action_2 = new Action(new Put(Bytes.toBytes("ccc")), 10);
    Action action_3 = new Action(new Delete(Bytes.toBytes("ccc")), 10);
    assertFalse(action_0.equals(action_1));
    assertTrue(action_0.equals(action_0));
    assertTrue(action_1.equals(action_2));
    assertTrue(action_2.equals(action_1));
    assertFalse(action_0.equals(new Put(Bytes.toBytes("abc"))));
    assertTrue(action_2.equals(action_3));
    assertFalse(action_0.equals(action_3));
    assertEquals(0, action_0.compareTo(action_0));
    assertTrue(action_0.compareTo(action_1) < 0);
    assertTrue(action_1.compareTo(action_0) > 0);
    assertEquals(0, action_1.compareTo(action_2));
  }

  @Test
  public void testBatch() throws IOException, InterruptedException {
    ClusterConnection conn = new MyConnectionImpl(CONF);
    HTable ht = (HTable) conn.getTable(DUMMY_TABLE);
    ht.multiAp = new MyAsyncProcess(conn, CONF);

    List<Put> puts = new ArrayList<>(7);
    puts.add(createPut(1, true));
    puts.add(createPut(1, true));
    puts.add(createPut(1, true));
    puts.add(createPut(1, true));
    puts.add(createPut(1, false)); // <=== the bad apple, position 4
    puts.add(createPut(1, true));
    puts.add(createPut(1, false)); // <=== another bad apple, position 6

    Object[] res = new Object[puts.size()];
    try {
      ht.batch(puts, res);
      Assert.fail();
    } catch (RetriesExhaustedException expected) {
    }

    Assert.assertEquals(success, res[0]);
    Assert.assertEquals(success, res[1]);
    Assert.assertEquals(success, res[2]);
    Assert.assertEquals(success, res[3]);
    Assert.assertEquals(failure, res[4]);
    Assert.assertEquals(success, res[5]);
    Assert.assertEquals(failure, res[6]);
  }

  @Test
  public void testErrorsServers() throws IOException {
    Configuration configuration = new Configuration(CONF);
    ClusterConnection conn = new MyConnectionImpl(configuration);
    MyAsyncProcess ap = new MyAsyncProcess(conn, configuration);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);
    configuration.setBoolean(ConnectionImplementation.RETRIES_BY_SERVER_KEY, true);

    Assert.assertNotNull(ap.createServerErrorTracker());
    Assert.assertTrue(ap.serverTrackerTimeout > 200L);
    ap.serverTrackerTimeout = 1L;

    Put p = createPut(1, false);
    mutator.mutate(p);

    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    // Checking that the ErrorsServers came into play and didn't make us stop immediately
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());
  }

  @Test
  public void testReadAndWriteTimeout() throws IOException {
    final long readTimeout = 10 * 1000;
    final long writeTimeout = 20 * 1000;
    Configuration copyConf = new Configuration(CONF);
    copyConf.setLong(HConstants.HBASE_RPC_READ_TIMEOUT_KEY, readTimeout);
    copyConf.setLong(HConstants.HBASE_RPC_WRITE_TIMEOUT_KEY, writeTimeout);
    ClusterConnection conn = new MyConnectionImpl(copyConf);
    MyAsyncProcess ap = new MyAsyncProcess(conn, copyConf);
    try (HTable ht = (HTable) conn.getTable(DUMMY_TABLE)) {
      ht.multiAp = ap;
      List<Get> gets = new ArrayList<>();
      gets.add(new Get(DUMMY_BYTES_1));
      gets.add(new Get(DUMMY_BYTES_2));
      try {
        ht.get(gets);
      } catch (ClassCastException e) {
        // No result response on this test.
      }
      assertEquals(readTimeout, ap.previousTimeout);
      ap.previousTimeout = -1;

      try {
        ht.existsAll(gets);
      } catch (ClassCastException e) {
        // No result response on this test.
      }
      assertEquals(readTimeout, ap.previousTimeout);
      ap.previousTimeout = -1;

      List<Delete> deletes = new ArrayList<>();
      deletes.add(new Delete(DUMMY_BYTES_1));
      deletes.add(new Delete(DUMMY_BYTES_2));
      ht.delete(deletes);
      assertEquals(writeTimeout, ap.previousTimeout);
    }
  }

  @Test
  public void testErrors() throws IOException {
    ClusterConnection conn = new MyConnectionImpl(CONF);
    AsyncProcessWithFailure ap = new AsyncProcessWithFailure(conn, CONF, new IOException("test"));
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);

    Assert.assertNotNull(ap.createServerErrorTracker());

    Put p = createPut(1, true);
    mutator.mutate(p);

    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    // Checking that the ErrorsServers came into play and didn't make us stop immediately
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());
  }

  @Test
  public void testCallQueueTooLarge() throws IOException {
    ClusterConnection conn = new MyConnectionImpl(CONF);
    AsyncProcessWithFailure ap =
      new AsyncProcessWithFailure(conn, CONF, new CallQueueTooBigException());
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);
    Assert.assertNotNull(ap.createServerErrorTracker());
    Put p = createPut(1, true);
    mutator.mutate(p);

    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    // Checking that the ErrorsServers came into play and didn't make us stop immediately
    Assert.assertEquals(NB_RETRIES + 1, ap.callsCt.get());
  }

  /**
   * This test simulates multiple regions on 2 servers. We should have 2 multi requests and 2
   * threads: 1 per server, this whatever the number of regions.
   */
  @Test
  public void testThreadCreation() throws Exception {
    final int NB_REGS = 100;
    List<HRegionLocation> hrls = new ArrayList<>(NB_REGS);
    List<Get> gets = new ArrayList<>(NB_REGS);
    for (int i = 0; i < NB_REGS; i++) {
      HRegionInfo hri =
        new HRegionInfo(DUMMY_TABLE, Bytes.toBytes(i * 10L), Bytes.toBytes(i * 10L + 9L), false, i);
      HRegionLocation hrl = new HRegionLocation(hri, i % 2 == 0 ? sn : sn2);
      hrls.add(hrl);

      Get get = new Get(Bytes.toBytes(i * 10L));
      gets.add(get);
    }

    MyConnectionImpl2 con = new MyConnectionImpl2(hrls, CONF);
    MyAsyncProcess ap = new MyAsyncProcess(con, CONF, con.nbThreads);
    HTable ht = (HTable) con.getTable(DUMMY_TABLE, ap.service);
    ht.multiAp = ap;
    ht.batch(gets, null);

    Assert.assertEquals(NB_REGS, ap.nbActions.get());
    Assert.assertEquals("1 multi response per server", 2, ap.nbMultiResponse.get());
    Assert.assertEquals("1 thread per server", 2, con.nbThreads.get());

    int nbReg = 0;
    for (int i = 0; i < NB_REGS; i++) {
      if (con.usedRegions[i]) nbReg++;
    }
    Assert.assertEquals("nbReg=" + nbReg, NB_REGS, nbReg);
  }

  @Test
  public void testReplicaReplicaSuccess() throws Exception {
    // Main call takes too long so replicas succeed, except for one region w/o replicas.
    // One region has no replica, so the main call succeeds for it.
    MyAsyncProcessWithReplicas ap = createReplicaAp(10, 1000, 0);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2, DUMMY_BYTES_3);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[3]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.TRUE, RR.TRUE, RR.FALSE);
    Assert.assertEquals(2, ap.getReplicaCallCount());
  }

  @Test
  public void testReplicaPrimarySuccessWoReplicaCalls() throws Exception {
    // Main call succeeds before replica calls are kicked off.
    MyAsyncProcessWithReplicas ap = createReplicaAp(1000, 10, 0);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2, DUMMY_BYTES_3);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[3]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.FALSE, RR.FALSE, RR.FALSE);
    Assert.assertEquals(0, ap.getReplicaCallCount());
  }

  @Test
  public void testReplicaParallelCallsSucceed() throws Exception {
    // Either main or replica can succeed.
    MyAsyncProcessWithReplicas ap = createReplicaAp(0, 0, 0);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[2]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.DONT_CARE, RR.DONT_CARE);
    long replicaCalls = ap.getReplicaCallCount();
    Assert.assertTrue(replicaCalls >= 0);
    Assert.assertTrue(replicaCalls <= 2);
  }

  @Test
  public void testReplicaPartialReplicaCall() throws Exception {
    // One server is slow, so the result for its region comes from replica, whereas
    // the result for other region comes from primary before replica calls happen.
    // There should be no replica call for that region at all.
    MyAsyncProcessWithReplicas ap = createReplicaAp(1000, 0, 0);
    ap.setPrimaryCallDelay(sn2, 2000);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[2]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.FALSE, RR.TRUE);
    Assert.assertEquals(1, ap.getReplicaCallCount());
  }

  @Test
  public void testReplicaMainFailsBeforeReplicaCalls() throws Exception {
    // Main calls fail before replica calls can start - this is currently not handled.
    // It would probably never happen if we can get location (due to retries),
    // and it would require additional synchronization.
    MyAsyncProcessWithReplicas ap = createReplicaAp(1000, 0, 0, 0);
    ap.addFailures(hri1, hri2);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[2]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.FAILED, RR.FAILED);
    Assert.assertEquals(0, ap.getReplicaCallCount());
  }

  @Test
  public void testReplicaReplicaSuccessWithParallelFailures() throws Exception {
    // Main calls fails after replica calls start. For two-replica region, one replica call
    // also fails. Regardless, we get replica results for both regions.
    MyAsyncProcessWithReplicas ap = createReplicaAp(0, 1000, 1000, 0);
    ap.addFailures(hri1, hri1r2, hri2);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[2]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.TRUE, RR.TRUE);
    Assert.assertEquals(2, ap.getReplicaCallCount());
  }

  @Test
  public void testReplicaAllCallsFailForOneRegion() throws Exception {
    // For one of the region, all 3, main and replica, calls fail. For the other, replica
    // call fails but its exception should not be visible as it did succeed.
    MyAsyncProcessWithReplicas ap = createReplicaAp(500, 1000, 0, 0);
    ap.addFailures(hri1, hri1r1, hri1r2, hri2r1);
    List<Get> rows = makeTimelineGets(DUMMY_BYTES_1, DUMMY_BYTES_2);
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(ap.service)
      .setRpcTimeout(RPC_TIMEOUT).setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE)
      .setRowAccess(rows).setResults(new Object[2]).setSubmittedRows(SubmittedRows.ALL).build();
    AsyncRequestFuture ars = ap.submit(task);
    verifyReplicaResult(ars, RR.FAILED, RR.FALSE);
    // We should get 3 exceptions, for main + 2 replicas for DUMMY_BYTES_1
    Assert.assertEquals(3, ars.getErrors().getNumExceptions());
    for (int i = 0; i < ars.getErrors().getNumExceptions(); ++i) {
      Assert.assertArrayEquals(DUMMY_BYTES_1, ars.getErrors().getRow(i).getRow());
    }
  }

  private MyAsyncProcessWithReplicas createReplicaAp(int replicaAfterMs, int primaryMs,
    int replicaMs) throws Exception {
    return createReplicaAp(replicaAfterMs, primaryMs, replicaMs, -1);
  }

  private MyAsyncProcessWithReplicas createReplicaAp(int replicaAfterMs, int primaryMs,
    int replicaMs, int retries) throws Exception {
    // TODO: this is kind of timing dependent... perhaps it should detect from createCaller
    // that the replica call has happened and that way control the ordering.
    Configuration conf = new Configuration();
    conf.setInt(AsyncProcess.PRIMARY_CALL_TIMEOUT_KEY, replicaAfterMs * 1000);
    if (retries >= 0) {
      conf.setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, retries);
    }
    ClusterConnection conn = createHConnectionWithReplicas(new ConnectionConfiguration(conf));
    MyAsyncProcessWithReplicas ap = new MyAsyncProcessWithReplicas(conn, conf);
    ap.setCallDelays(primaryMs, replicaMs);
    return ap;
  }

  private BufferedMutatorParams createBufferedMutatorParams(MyAsyncProcess ap, TableName name) {
    return new BufferedMutatorParams(name).pool(ap.service).rpcTimeout(RPC_TIMEOUT)
      .opertationTimeout(OPERATION_TIMEOUT);
  }

  private static List<Get> makeTimelineGets(byte[]... rows) {
    List<Get> result = new ArrayList<>(rows.length);
    for (byte[] row : rows) {
      Get get = new Get(row);
      get.setConsistency(Consistency.TIMELINE);
      result.add(get);
    }
    return result;
  }

  private void verifyResult(AsyncRequestFuture ars, boolean... expected) throws Exception {
    Object[] actual = ars.getResults();
    Assert.assertEquals(expected.length, actual.length);
    for (int i = 0; i < expected.length; ++i) {
      Assert.assertEquals(expected[i], !(actual[i] instanceof Throwable));
    }
  }

  /** After reading TheDailyWtf, I always wanted to create a MyBoolean enum like this! */
  private enum RR {
    TRUE,
    FALSE,
    DONT_CARE,
    FAILED
  }

  private void verifyReplicaResult(AsyncRequestFuture ars, RR... expecteds) throws Exception {
    Object[] actuals = ars.getResults();
    Assert.assertEquals(expecteds.length, actuals.length);
    for (int i = 0; i < expecteds.length; ++i) {
      Object actual = actuals[i];
      RR expected = expecteds[i];
      Assert.assertEquals(actual.toString(), expected == RR.FAILED, actual instanceof Throwable);
      if (expected != RR.FAILED && expected != RR.DONT_CARE) {
        Assert.assertEquals(expected == RR.TRUE, ((Result) actual).isStale());
      }
    }
  }

  /**
   * @param regCnt  the region: 1 to 3.
   * @param success if true, the put will succeed.
   * @return a put
   */
  private Put createPut(int regCnt, boolean success) {
    Put p;
    if (!success) {
      p = new Put(FAILS);
    } else switch (regCnt) {
      case 1:
        p = new Put(DUMMY_BYTES_1);
        break;
      case 2:
        p = new Put(DUMMY_BYTES_2);
        break;
      case 3:
        p = new Put(DUMMY_BYTES_3);
        break;
      default:
        throw new IllegalArgumentException("unknown " + regCnt);
    }

    p.addColumn(DUMMY_BYTES_1, DUMMY_BYTES_1, DUMMY_BYTES_1);

    return p;
  }

  static class MyThreadPoolExecutor extends ThreadPoolExecutor {
    public MyThreadPoolExecutor(int coreThreads, int maxThreads, long keepAliveTime,
      TimeUnit timeunit, BlockingQueue<Runnable> blockingqueue) {
      super(coreThreads, maxThreads, keepAliveTime, timeunit, blockingqueue);
    }

    @Override
    public Future submit(Runnable runnable) {
      throw new OutOfMemoryError("OutOfMemory error thrown by means");
    }
  }

  static class AsyncProcessForThrowableCheck extends AsyncProcess {
    public AsyncProcessForThrowableCheck(ClusterConnection hc, Configuration conf) {
      super(hc, conf, new RpcRetryingCallerFactory(conf, hc.getConnectionConfiguration()),
        new RpcControllerFactory(conf));
    }
  }

  @Test
  public void testUncheckedException() throws Exception {
    // Test the case pool.submit throws unchecked exception
    ClusterConnection hc = createHConnection();
    MyThreadPoolExecutor myPool =
      new MyThreadPoolExecutor(1, 20, 60, TimeUnit.SECONDS, new LinkedBlockingQueue<>(200));
    AsyncProcess ap = new AsyncProcessForThrowableCheck(hc, CONF);

    List<Put> puts = new ArrayList<>(1);
    puts.add(createPut(1, true));
    AsyncProcessTask task = AsyncProcessTask.newBuilder().setPool(myPool).setRpcTimeout(RPC_TIMEOUT)
      .setOperationTimeout(OPERATION_TIMEOUT).setTableName(DUMMY_TABLE).setRowAccess(puts)
      .setSubmittedRows(SubmittedRows.NORMAL).build();
    ap.submit(task);
    Assert.assertTrue(puts.isEmpty());
  }

  /**
   * Below tests make sure we could use a special pause setting when retry an exception where
   * {@link HBaseServerException#isServerOverloaded(Throwable)} is true, see HBASE-17114
   */

  @Test
  public void testRetryPauseWhenServerOverloadedDueToCQTBE() throws Exception {
    testRetryPauseWhenServerIsOverloaded(new CallQueueTooBigException());
  }

  @Test
  public void testRetryPauseWhenServerOverloadedDueToCDE() throws Exception {
    testRetryPauseWhenServerIsOverloaded(new CallDroppedException());
  }

  @Test
  public void testRetryPauseForRpcThrottling() throws IOException {
    long waitInterval = 500L;
    testRetryPause(new Configuration(CONF), waitInterval, new RpcThrottlingException(
      RpcThrottlingException.Type.NumReadRequestsExceeded, waitInterval, "For test"));
  }

  private void testRetryPauseWhenServerIsOverloaded(HBaseServerException exception)
    throws IOException {
    Configuration testConf = new Configuration(CONF);
    long specialPause = 500L;
    testConf.setLong(ConnectionConfiguration.HBASE_CLIENT_PAUSE_FOR_SERVER_OVERLOADED,
      specialPause);
    testRetryPause(testConf, specialPause, exception);
  }

  private void testRetryPause(Configuration testConf, long expectedPause,
    HBaseIOException exception) throws IOException {

    final int retries = 1;
    testConf.setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, retries);

    ClusterConnection conn = new MyConnectionImpl(testConf);
    AsyncProcessWithFailure ap = new AsyncProcessWithFailure(conn, testConf, exception);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);

    Assert.assertNotNull(mutator.getAsyncProcess().createServerErrorTracker());

    Put p = createPut(1, true);
    mutator.mutate(p);

    long startTime = EnvironmentEdgeManager.currentTime();
    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    long actualSleep = EnvironmentEdgeManager.currentTime() - startTime;
    long expectedSleep = 0L;
    for (int i = 0; i < retries; i++) {
      expectedSleep += ConnectionUtils.getPauseTime(expectedPause, i);
      // Prevent jitter in ConcurrentMapUtils#getPauseTime to affect result
      actualSleep += (long) (expectedPause * 0.01f);
    }
    LOG.debug("Expected to sleep " + expectedSleep + "ms, actually slept " + actualSleep + "ms");
    Assert.assertTrue("Expected to sleep " + expectedSleep + " but actually " + actualSleep + "ms",
      actualSleep >= expectedSleep);

    // check and confirm normal IOE will use the normal pause
    final long normalPause =
      testConf.getLong(HConstants.HBASE_CLIENT_PAUSE, HConstants.DEFAULT_HBASE_CLIENT_PAUSE);
    ap = new AsyncProcessWithFailure(conn, testConf, new IOException());
    bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    mutator = new BufferedMutatorImpl(conn, bufferParam, ap);
    Assert.assertNotNull(mutator.getAsyncProcess().createServerErrorTracker());
    mutator.mutate(p);
    startTime = EnvironmentEdgeManager.currentTime();
    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getRow(0) == p);
    }
    actualSleep = EnvironmentEdgeManager.currentTime() - startTime;
    expectedSleep = 0L;
    for (int i = 0; i < retries; i++) {
      expectedSleep += ConnectionUtils.getPauseTime(normalPause, i);
    }
    // plus an additional pause to balance the program execution time
    expectedSleep += normalPause;
    LOG.debug("Expected to sleep " + expectedSleep + "ms, actually slept " + actualSleep + "ms");
    Assert.assertTrue("Slept for too long: " + actualSleep + "ms", actualSleep <= expectedSleep);
  }

  @Test
  public void testFastFailIfBackoffGreaterThanRemaining() throws IOException {
    Configuration testConf = new Configuration(CONF);
    testConf.setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 100);
    long waitInterval = 500L;
    HBaseIOException exception = new RpcThrottlingException(
      RpcThrottlingException.Type.NumReadRequestsExceeded, waitInterval, "For test");

    final int retries = 1;
    testConf.setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, retries);

    ClusterConnection conn = new MyConnectionImpl(testConf);
    AsyncProcessWithFailure ap = new AsyncProcessWithFailure(conn, testConf, exception);
    BufferedMutatorParams bufferParam =
      createBufferedMutatorParams(ap, DUMMY_TABLE).operationTimeout(100);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);

    Assert.assertNotNull(mutator.getAsyncProcess().createServerErrorTracker());

    Put p = createPut(1, true);
    mutator.mutate(p);

    try {
      mutator.flush();
      Assert.fail();
    } catch (RetriesExhaustedWithDetailsException expected) {
      assertEquals(1, expected.getNumExceptions());
      assertTrue(expected.getCause(0) instanceof OperationTimeoutExceededException);
      assertTrue(expected.getCause(0).getMessage().startsWith("Backoff"));
    }
  }

  /**
   * Tests that we properly recover from exceptions that DO NOT go through receiveGlobalFailure, due
   * to updating the meta cache for the region which failed. Successful multigets can include region
   * exceptions in the MultiResponse. In that case, it skips receiveGlobalFailure and instead
   * handles in receiveMultiAction
   */
  @Test
  public void testRetryWithExceptionClearsMetaCacheUsingRegionException() throws Exception {
    testRetryWithExceptionClearsMetaCache(true);
  }

  /**
   * Tests that we properly recover from exceptions that go through receiveGlobalFailure, due to
   * updating the meta cache for the region which failed.
   */
  @Test
  public void testRetryWithExceptionClearsMetaCacheUsingServerException() throws Exception {
    testRetryWithExceptionClearsMetaCache(false);
  }

  private void testRetryWithExceptionClearsMetaCache(boolean useRegionException)
    throws IOException {
    Configuration myConf = new Configuration(CONF);
    myConf.setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 1);
    ClusterConnection conn = createHConnection(new ConnectionConfiguration(myConf));

    // we pass in loc1.getServerName here so that only calls to that server will fail
    AsyncProcessWithFailure ap = new AsyncProcessWithFailure(conn, myConf,
      new RegionOpeningException("test"), loc1.getServerName(), useRegionException);
    BufferedMutatorParams bufferParam = createBufferedMutatorParams(ap, DUMMY_TABLE);
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, bufferParam, ap);

    Assert.assertNotNull(mutator.getAsyncProcess().createServerErrorTracker());

    Assert.assertEquals(conn.locateRegion(DUMMY_TABLE, DUMMY_BYTES_1, true, true).toString(),
      new RegionLocations(loc1).toString());

    // simulate updateCachedLocations, by changing the loc for this row to loc3. only loc1 fails,
    // so this means retry will succeed
    Mockito.doAnswer(invocation -> {
      setMockLocation(conn, DUMMY_BYTES_1, new RegionLocations(loc3));
      return null;
    }).when(conn).updateCachedLocations(Mockito.eq(DUMMY_TABLE),
      Mockito.eq(loc1.getRegion().getRegionName()), Mockito.eq(DUMMY_BYTES_1), Mockito.any(),
      Mockito.eq(loc1.getServerName()));

    // Ensure we haven't called updateCachedLocations yet
    Mockito.verify(conn, Mockito.times(0)).updateCachedLocations(Mockito.any(), Mockito.any(),
      Mockito.any(), Mockito.any(), Mockito.any());

    Put p = createPut(1, true);
    mutator.mutate(p);

    // we expect this to succeed because the bad region location should be updated upon
    // the initial failure causing retries to succeed.
    mutator.flush();

    // validate that we updated the location, as we expected
    Assert.assertEquals(conn.locateRegion(DUMMY_TABLE, DUMMY_BYTES_1, true, true).toString(),
      new RegionLocations(loc3).toString());
    // this is a given since the location updated, but validate that we called updateCachedLocations
    Mockito.verify(conn, Mockito.atLeastOnce()).updateCachedLocations(Mockito.eq(DUMMY_TABLE),
      Mockito.eq(loc1.getRegion().getRegionName()), Mockito.eq(DUMMY_BYTES_1), Mockito.any(),
      Mockito.eq(loc1.getServerName()));
  }

  @Test
  public void testQueueRowAccess() throws Exception {
    ClusterConnection conn = createHConnection();
    BufferedMutatorImpl mutator = new BufferedMutatorImpl(conn, null, null,
      new BufferedMutatorParams(DUMMY_TABLE).writeBufferSize(100000));
    Put p0 = new Put(DUMMY_BYTES_1).addColumn(DUMMY_BYTES_1, DUMMY_BYTES_1, DUMMY_BYTES_1);
    Put p1 = new Put(DUMMY_BYTES_2).addColumn(DUMMY_BYTES_2, DUMMY_BYTES_2, DUMMY_BYTES_2);
    mutator.mutate(p0);
    BufferedMutatorImpl.QueueRowAccess ra0 = mutator.createQueueRowAccess();
    // QueueRowAccess should take all undealt mutations
    assertEquals(0, mutator.size());
    mutator.mutate(p1);
    assertEquals(1, mutator.size());
    BufferedMutatorImpl.QueueRowAccess ra1 = mutator.createQueueRowAccess();
    // QueueRowAccess should take all undealt mutations
    assertEquals(0, mutator.size());
    assertEquals(1, ra0.size());
    assertEquals(1, ra1.size());
    Iterator<Row> iter0 = ra0.iterator();
    Iterator<Row> iter1 = ra1.iterator();
    assertTrue(iter0.hasNext());
    assertTrue(iter1.hasNext());
    // the next() will poll the mutation from inner buffer and update the buffer count
    assertTrue(iter0.next() == p0);
    assertEquals(1, mutator.getUnflushedSize());
    assertEquals(p1.heapSize(), mutator.getCurrentWriteBufferSize());
    assertTrue(iter1.next() == p1);
    assertEquals(0, mutator.getUnflushedSize());
    assertEquals(0, mutator.getCurrentWriteBufferSize());
    assertFalse(iter0.hasNext());
    assertFalse(iter1.hasNext());
    // ra0 doest handle the mutation so the mutation won't be pushed back to buffer
    iter0.remove();
    ra0.close();
    assertEquals(0, mutator.size());
    assertEquals(0, mutator.getUnflushedSize());
    assertEquals(0, mutator.getCurrentWriteBufferSize());
    // ra1 doesn't handle the mutation so the mutation will be pushed back to buffer
    ra1.close();
    assertEquals(1, mutator.size());
    assertEquals(1, mutator.getUnflushedSize());
    assertEquals(p1.heapSize(), mutator.getCurrentWriteBufferSize());
  }
}