release v3.1.0

Signed-off-by: birdstorm <samuelwyf@hotmail.com>

.ci/integration_test.groovy

@@ -9,6 +9,7 @@ def call(ghprbActualCommit, ghprbPullId, ghprbPullTitle, ghprbPullLink, ghprbPul
     if (m1) {
         TIDB_BRANCH = "${m1[0][1]}"
     }
+    m1 = null
     println "TIDB_BRANCH=${TIDB_BRANCH}"
 
     // parse pd branch
@@ -16,6 +17,7 @@ def call(ghprbActualCommit, ghprbPullId, ghprbPullTitle, ghprbPullLink, ghprbPul
     if (m2) {
         PD_BRANCH = "${m2[0][1]}"
     }
+    m2 = null
     println "PD_BRANCH=${PD_BRANCH}"
 
     // parse tikv branch
@@ -23,6 +25,7 @@ def call(ghprbActualCommit, ghprbPullId, ghprbPullTitle, ghprbPullLink, ghprbPul
     if (m3) {
         TIKV_BRANCH = "${m3[0][1]}"
     }
+    m3 = null
     println "TIKV_BRANCH=${TIKV_BRANCH}"
 
     catchError {
@@ -61,13 +64,13 @@ def call(ghprbActualCommit, ghprbPullId, ghprbPullTitle, ghprbPullLink, ghprbPul
                         killall -9 pd-server || true
                         killall -9 java || true
                         sleep 10
-                        bin/pd-server --name=pd --data-dir=pd --config=../.ci/config/pd.toml &>pd.log &
+                        bin/pd-server --name=pd --data-dir=pd --config=../config/pd.toml &>pd.log &
                         sleep 10
-                        bin/tikv-server --pd=127.0.0.1:2379 -s tikv --addr=0.0.0.0:20160 --advertise-addr=127.0.0.1:20160 --config=../.ci/config/tikv.toml &>tikv.log &
+                        bin/tikv-server --pd=127.0.0.1:2379 -s tikv --addr=0.0.0.0:20160 --advertise-addr=127.0.0.1:20160 --config=../config/tikv.toml &>tikv.log &
                         sleep 10
                         ps aux | grep '-server' || true
                         curl -s 127.0.0.1:2379/pd/api/v1/status || true
-                        bin/tidb-server --store=tikv --path="127.0.0.1:2379" --config=../.ci/config/tidb.toml &>tidb.log &
+                        bin/tidb-server --store=tikv --path="127.0.0.1:2379" --config=../config/tidb.toml &>tidb.log &
                         sleep 60
                         """
                     }

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,36 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: "[BUG] Title of Bug Report"
+labels: type/bug
+assignees: ''
+
+---
+
+**Describe the bug**
+<!-- A clear and concise description of what the bug is. -->
+
+**What did you do**
+<!--
+If possible, please provide a code receipt to produce this issue.
+ -->
+
+**What do you expect**
+<!-- A clear and concise description of what you expected to happen. -->
+
+**What happens instead**
+<!-- If an error occurs, please provide complete error stack. -->
+
+<!--
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+ -->
+
+**Java Client and TiDB/TiKV version info**
+<!-- What version of Spark and TiSpark are you using? (Provide Spark version and run `spark.sql(“select ti_version()”).show(false)` in spark-shell) -->
+
+<!--
+**Additional context**
+Add any other context about the problem here.
+You may also provide TiDB version here if it is related to the issue.
+-->

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: type/feature-request
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

.github/pull_request_template.md

@@ -0,0 +1,33 @@
+### What problem does this PR solve? <!--add issue link with summary if exists-->
+
+
+### What is changed and how it works?
+
+
+### Check List <!--REMOVE the items that are not applicable-->
+
+Tests <!-- At least one of them must be included. -->
+
+ - Unit test
+ - Integration test
+ - Manual test (add detailed scripts or steps below)
+ - No code
+
+Code changes
+
+ - Has exported function/method change
+ - Has exported variable/fields change
+ - Has interface methods change
+ - Has persistent data change
+
+Side effects
+
+ - Possible performance regression
+ - Increased code complexity
+ - Breaking backward compatibility
+
+Related changes
+
+ - Need to cherry-pick to the release branch
+ - Need to update the documentation
+ - Need to be included in the release note

config/pd.toml

@@ -0,0 +1,4 @@
+# PD Configuration.
+[replication]
+enable-placement-rules = true
+max-replicas = 1

config/tidb.toml

@@ -0,0 +1 @@
+# TiDB Configuration.

config/tikv.toml

@@ -0,0 +1,5 @@
+# TiKV Configuration.
+
+[raftstore]
+# set store capacity, if no set, use disk capacity.
+capacity = "8G"

pom.xml

@@ -5,7 +5,7 @@
 
     <groupId>org.tikv</groupId>
     <artifactId>tikv-client-java</artifactId>
-    <version>3.0.2-SNAPSHOT</version>
+    <version>3.1.0</version>
     <packaging>jar</packaging>
     <name>TiKV Java Client</name>
     <description>A Java Client for TiKV</description>
@@ -126,6 +126,11 @@
             <artifactId>grpc-stub</artifactId>
             <version>${grpc.version}</version>
         </dependency>
+        <dependency>
+            <groupId>io.grpc</groupId>
+            <artifactId>grpc-services</artifactId>
+            <version>${grpc.version}</version>
+        </dependency>
         <dependency>
             <groupId>io.grpc</groupId>
             <artifactId>grpc-testing</artifactId>

src/main/java/org/tikv/BUILD

@@ -1,34 +0,0 @@
-package(default_visibility = ["//visibility:public"])
-
-java_library(
-    name = "tikv-java-client-lib",
-    srcs = glob(
-        ["**/*.java"],
-    ),
-    deps = [
-        "//:java",
-        "@com_fasterxml_jackson_core_jackson_annotations//jar",
-        "@com_fasterxml_jackson_core_jackson_core//jar",
-        "@com_fasterxml_jackson_core_jackson_databind//jar",
-        "@com_google_code_findbugs_jsr305//jar",
-        "@com_google_code_gson_gson//jar",
-        "@com_google_errorprone_error_prone_annotations//jar",
-        "@com_google_guava_guava//jar",
-        "@com_google_protobuf_protobuf_java//jar",
-        "@joda_time//jar",
-        # the following are defined in rules_protobuf
-        "@org_pubref_rules_protobuf//java:grpc_compiletime_deps",
-        "@org_pubref_rules_protobuf//java:netty_runtime_deps",
-
-        "@org_slf4j_slf4j_api//jar",
-        "@org_slf4j_jcl_over_slf4j//jar",
-        "@org_slf4j_jul_to_slf4j//jar",
-        "@log4j_log4j//jar",
-        "@net_sf_trove4j_trove4j//jar",
-    ],
-)
-
-filegroup(
-    name = "srcs",
-    srcs = ["BUILD"] + glob(["**/*.java"]),
-)

src/main/java/org/tikv/Main.java

@@ -1,5 +0,0 @@
-package org.tikv;
-
-public class Main {
-  public static void main(String args[]) throws Exception {}
-}

src/main/java/org/tikv/common/AbstractGRPCClient.java

@@ -18,10 +18,15 @@ package org.tikv.common;
 import static io.grpc.stub.ClientCalls.asyncBidiStreamingCall;
 import static io.grpc.stub.ClientCalls.blockingServerStreamingCall;
 
+import io.grpc.ManagedChannel;
 import io.grpc.MethodDescriptor;
+import io.grpc.health.v1.HealthCheckRequest;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
 import io.grpc.stub.AbstractStub;
 import io.grpc.stub.ClientCalls;
 import io.grpc.stub.StreamObserver;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -171,4 +176,20 @@ public abstract class AbstractGRPCClient<
   protected abstract BlockingStubT getBlockingStub();
 
   protected abstract StubT getAsyncStub();
+
+  protected boolean checkHealth(String addressStr, HostMapping hostMapping) {
+    ManagedChannel channel = channelFactory.getChannel(addressStr, hostMapping);
+    HealthGrpc.HealthBlockingStub stub =
+        HealthGrpc.newBlockingStub(channel).withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
+    HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
+    try {
+      HealthCheckResponse resp = stub.check(req);
+      if (resp.getStatus() != HealthCheckResponse.ServingStatus.SERVING) {
+        return false;
+      }
+    } catch (Exception e) {
+      return false;
+    }
+    return true;
+  }
 }

src/main/java/org/tikv/common/ConfigUtils.java

@@ -42,16 +42,19 @@ public class ConfigUtils {
   public static final String TIKV_KV_CLIENT_CONCURRENCY = "tikv.kv_client_concurrency";
 
   public static final String TIKV_KV_MODE = "tikv.kv_mode";
-  public static final String TIKV_IS_REPLICA_READ = "tikv.is_replica_read";
+  public static final String TIKV_REPLICA_READ = "tikv.replica_read";
 
   public static final String TIKV_METRICS_ENABLE = "tikv.metrics.enable";
   public static final String TIKV_METRICS_PORT = "tikv.metrics.port";
 
   public static final String TIKV_NETWORK_MAPPING_NAME = "tikv.network.mapping";
+  public static final String TIKV_ENABLE_GRPC_FORWARD = "tikv.enable_grpc_forward";
+  public static final String TIKV_GRPC_HEALTH_CHECK_TIMEOUT = "tikv.grpc.health_check_timeout";
 
   public static final String DEF_PD_ADDRESSES = "127.0.0.1:2379";
   public static final String DEF_TIMEOUT = "600ms";
   public static final String DEF_SCAN_TIMEOUT = "20s";
+  public static final int DEF_CHECK_HEALTH_TIMEOUT = 40;
   public static final int DEF_SCAN_BATCH_SIZE = 10240;
   public static final int DEF_MAX_FRAME_SIZE = 268435456 * 2; // 256 * 2 MB
   public static final int DEF_INDEX_SCAN_BATCH_SIZE = 20000;
@@ -72,10 +75,11 @@ public class ConfigUtils {
   public static final String DEF_DB_PREFIX = "";
   public static final int DEF_KV_CLIENT_CONCURRENCY = 10;
   public static final TiConfiguration.KVMode DEF_KV_MODE = TiConfiguration.KVMode.TXN;
-  public static final boolean DEF_IS_REPLICA_READ = false;
+  public static final String DEF_REPLICA_READ = "LEADER";
   public static final boolean DEF_METRICS_ENABLE = false;
   public static final int DEF_METRICS_PORT = 3140;
   public static final String DEF_TIKV_NETWORK_MAPPING_NAME = "";
+  public static final boolean DEF_GRPC_FORWARD_ENABLE = true;
 
   public static final String NORMAL_COMMAND_PRIORITY = "NORMAL";
   public static final String LOW_COMMAND_PRIORITY = "LOW";
@@ -86,4 +90,8 @@ public class ConfigUtils {
 
   public static final String RAW_KV_MODE = "RAW";
   public static final String TXN_KV_MODE = "TXN";
+
+  public static final String LEADER = "LEADER";
+  public static final String FOLLOWER = "FOLLOWER";
+  public static final String LEADER_AND_FOLLOWER = "LEADER_AND_FOLLOWER";
 }

src/main/java/org/tikv/common/PDClient.java

@@ -30,6 +30,8 @@ import io.etcd.jetcd.KeyValue;
 import io.etcd.jetcd.kv.GetResponse;
 import io.etcd.jetcd.options.GetOption;
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.stub.MetadataUtils;
 import io.prometheus.client.Histogram;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
@@ -39,10 +41,10 @@ import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.TiConfiguration.KVMode;
@@ -59,11 +61,11 @@ import org.tikv.common.util.BackOffFunction.BackOffFuncType;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
 import org.tikv.common.util.ConcreteBackOffer;
-import org.tikv.common.util.FutureObserver;
 import org.tikv.kvproto.Metapb.Store;
 import org.tikv.kvproto.PDGrpc;
 import org.tikv.kvproto.PDGrpc.PDBlockingStub;
 import org.tikv.kvproto.PDGrpc.PDStub;
+import org.tikv.kvproto.Pdpb;
 import org.tikv.kvproto.Pdpb.Error;
 import org.tikv.kvproto.Pdpb.ErrorType;
 import org.tikv.kvproto.Pdpb.GetAllStoresRequest;
@@ -91,7 +93,7 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   private final Logger logger = LoggerFactory.getLogger(PDClient.class);
   private RequestHeader header;
   private TsoRequest tsoReq;
-  private volatile LeaderWrapper leaderWrapper;
+  private volatile PDClientWrapper pdClientWrapper;
   private ScheduledExecutorService service;
   private ScheduledExecutorService tiflashReplicaService;
   private List<URI> pdAddrs;
@@ -241,37 +243,16 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
       return new TiRegion(
           resp.getRegion(),
           resp.getLeader(),
+          null,
           conf.getIsolationLevel(),
           conf.getCommandPriority(),
           conf.getKvMode(),
-          conf.isReplicaRead());
+          conf.getReplicaSelector());
     } finally {
       requestTimer.observeDuration();
     }
   }
 
-  @Override
-  public Future<TiRegion> getRegionByKeyAsync(BackOffer backOffer, ByteString key) {
-    FutureObserver<TiRegion, GetRegionResponse> responseObserver =
-        new FutureObserver<>(
-            resp ->
-                new TiRegion(
-                    resp.getRegion(),
-                    resp.getLeader(),
-                    conf.getIsolationLevel(),
-                    conf.getCommandPriority(),
-                    conf.getKvMode(),
-                    conf.isReplicaRead()));
-    Supplier<GetRegionRequest> request =
-        () -> GetRegionRequest.newBuilder().setHeader(header).setRegionKey(key).build();
-
-    PDErrorHandler<GetRegionResponse> handler =
-        new PDErrorHandler<>(getRegionResponseErrorExtractor, this);
-
-    callAsyncWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, responseObserver, handler);
-    return responseObserver.getFuture();
-  }
-
   @Override
   public TiRegion getRegionByID(BackOffer backOffer, long id) {
     Supplier<GetRegionByIDRequest> request =
@@ -285,33 +266,11 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
     return new TiRegion(
         resp.getRegion(),
         resp.getLeader(),
+        null,
         conf.getIsolationLevel(),
         conf.getCommandPriority(),
         conf.getKvMode(),
-        conf.isReplicaRead());
+        conf.getReplicaSelector());
-  }
-
-  @Override
-  public Future<TiRegion> getRegionByIDAsync(BackOffer backOffer, long id) {
-    FutureObserver<TiRegion, GetRegionResponse> responseObserver =
-        new FutureObserver<>(
-            resp ->
-                new TiRegion(
-                    resp.getRegion(),
-                    resp.getLeader(),
-                    conf.getIsolationLevel(),
-                    conf.getCommandPriority(),
-                    conf.getKvMode(),
-                    conf.isReplicaRead()));
-
-    Supplier<GetRegionByIDRequest> request =
-        () -> GetRegionByIDRequest.newBuilder().setHeader(header).setRegionId(id).build();
-    PDErrorHandler<GetRegionResponse> handler =
-        new PDErrorHandler<>(getRegionResponseErrorExtractor, this);
-
-    callAsyncWithRetry(
-        backOffer, PDGrpc.getGetRegionByIDMethod(), request, responseObserver, handler);
-    return responseObserver.getFuture();
   }
 
   private Supplier<GetStoreRequest> buildGetStoreReq(long storeId) {
@@ -334,20 +293,6 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
         .getStore();
   }
 
-  @Override
-  public Future<Store> getStoreAsync(BackOffer backOffer, long storeId) {
-    FutureObserver<Store, GetStoreResponse> responseObserver =
-        new FutureObserver<>(GetStoreResponse::getStore);
-
-    callAsyncWithRetry(
-        backOffer,
-        PDGrpc.getGetStoreMethod(),
-        buildGetStoreReq(storeId),
-        responseObserver,
-        buildPDErrorHandler());
-    return responseObserver.getFuture();
-  }
-
   @Override
   public List<Store> getAllStores(BackOffer backOffer) {
     return callWithRetry(
@@ -361,8 +306,8 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   }
 
   @Override
-  public boolean isReplicaRead() {
+  public TiConfiguration.ReplicaRead getReplicaRead() {
-    return conf.isReplicaRead();
+    return conf.getReplicaRead();
   }
 
   @Override
@@ -385,8 +330,8 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   }
 
   @VisibleForTesting
-  LeaderWrapper getLeaderWrapper() {
+  PDClientWrapper getPdClientWrapper() {
-    return leaderWrapper;
+    return pdClientWrapper;
   }
 
   private GetMembersResponse getMembers(URI uri) {
@@ -407,50 +352,127 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
     return null;
   }
 
-  synchronized boolean switchLeader(List<String> leaderURLs) {
+  // return whether the leader has changed to target address `leaderUrlStr`.
-    if (leaderURLs.isEmpty()) return false;
+  synchronized boolean trySwitchLeader(String leaderUrlStr) {
-    String leaderUrlStr = leaderURLs.get(0);
+    if (pdClientWrapper != null) {
-    // TODO: Why not strip protocol info on server side since grpc does not need it
+      if (leaderUrlStr.equals(pdClientWrapper.getLeaderInfo())) {
-    if (leaderWrapper != null && leaderUrlStr.equals(leaderWrapper.getLeaderInfo())) {
+        // The message to leader is not forwarded by follower.
-      return true;
+        if (leaderUrlStr.equals(pdClientWrapper.getStoreAddress())) {
+          return true;
+        }
+      }
+      // If leader has transfered to another member, we can create another leaderwrapper.
     }
     // switch leader
-    return createLeaderWrapper(leaderUrlStr);
+    return createLeaderClientWrapper(leaderUrlStr);
   }
 
-  private boolean createLeaderWrapper(String leaderUrlStr) {
+  private synchronized boolean createLeaderClientWrapper(String leaderUrlStr) {
     try {
-      URI newLeader = addrToUri(leaderUrlStr);
-      leaderUrlStr = uriToAddr(newLeader);
-      if (leaderWrapper != null && leaderUrlStr.equals(leaderWrapper.getLeaderInfo())) {
-        return true;
-      }
-
       // create new Leader
       ManagedChannel clientChannel = channelFactory.getChannel(leaderUrlStr, hostMapping);
-      leaderWrapper =
+      pdClientWrapper =
-          new LeaderWrapper(
+          new PDClientWrapper(leaderUrlStr, leaderUrlStr, clientChannel, System.nanoTime());
-              leaderUrlStr,
-              PDGrpc.newBlockingStub(clientChannel),
-              PDGrpc.newStub(clientChannel),
-              System.nanoTime());
     } catch (IllegalArgumentException e) {
       logger.error("Error updating leader. " + leaderUrlStr, e);
       return false;
     }
-    logger.info(String.format("Switched to new leader: %s", leaderWrapper));
+    logger.info(String.format("Switched to new leader: %s", pdClientWrapper));
     return true;
   }
 
-  public void updateLeader() {
+  synchronized boolean createFollowerClientWrapper(String followerUrlStr, String leaderUrls) {
+    // TODO: Why not strip protocol info on server side since grpc does not need it
+
+    try {
+      if (!checkHealth(followerUrlStr, hostMapping)) {
+        return false;
+      }
+
+      // create new Leader
+      ManagedChannel channel = channelFactory.getChannel(followerUrlStr, hostMapping);
+      pdClientWrapper = new PDClientWrapper(leaderUrls, followerUrlStr, channel, System.nanoTime());
+    } catch (IllegalArgumentException e) {
+      logger.error("Error updating follower. " + followerUrlStr, e);
+      return false;
+    }
+    logger.info(String.format("Switched to new leader by follower forward: %s", pdClientWrapper));
+    return true;
+  }
+
+  public synchronized void updateLeaderOrforwardFollower() {
     for (URI url : this.pdAddrs) {
       // since resp is null, we need update leader's address by walking through all pd server.
       GetMembersResponse resp = getMembers(url);
       if (resp == null) {
         continue;
       }
+      if (resp.getLeader().getClientUrlsList().isEmpty()) {
+        continue;
+      }
+
+      String leaderUrlStr = resp.getLeader().getClientUrlsList().get(0);
+      leaderUrlStr = uriToAddr(addrToUri(leaderUrlStr));
+
       // if leader is switched, just return.
-      if (switchLeader(resp.getLeader().getClientUrlsList())) {
+      if (checkHealth(leaderUrlStr, hostMapping) && trySwitchLeader(leaderUrlStr)) {
+        return;
+      }
+
+      if (!conf.getEnableGrpcForward()) {
+        continue;
+      }
+
+      List<Pdpb.Member> members = resp.getMembersList();
+
+      boolean hasReachNextMember = false;
+      // If we have not used follower forward, try the first follower.
+      if (pdClientWrapper != null && pdClientWrapper.getStoreAddress().equals(leaderUrlStr)) {
+        hasReachNextMember = true;
+      }
+
+      for (int i = 0; i < members.size() * 2; i++) {
+        Pdpb.Member member = members.get(i % members.size());
+        if (member.getMemberId() == resp.getLeader().getMemberId()) {
+          continue;
+        }
+        String followerUrlStr = member.getClientUrlsList().get(0);
+        followerUrlStr = uriToAddr(addrToUri(followerUrlStr));
+        if (pdClientWrapper != null && pdClientWrapper.getStoreAddress().equals(followerUrlStr)) {
+          hasReachNextMember = true;
+          continue;
+        }
+        if (hasReachNextMember && createFollowerClientWrapper(followerUrlStr, leaderUrlStr)) {
+          return;
+        }
+      }
+    }
+    if (pdClientWrapper == null) {
+      throw new TiClientInternalException(
+          "already tried all address on file, but not leader found yet.");
+    }
+  }
+
+  public void tryUpdateLeader() {
+    for (URI url : this.pdAddrs) {
+      // since resp is null, we need update leader's address by walking through all pd server.
+      GetMembersResponse resp = getMembers(url);
+      if (resp == null) {
+        continue;
+      }
+      List<URI> urls =
+          resp.getMembersList()
+              .stream()
+              .map(mem -> addrToUri(mem.getClientUrls(0)))
+              .collect(Collectors.toList());
+      String leaderUrlStr = resp.getLeader().getClientUrlsList().get(0);
+      leaderUrlStr = uriToAddr(addrToUri(leaderUrlStr));
+
+      // If leader is not change but becomes available, we can cancel follower forward.
+      if (checkHealth(leaderUrlStr, hostMapping) && trySwitchLeader(leaderUrlStr)) {
+        if (!urls.equals(this.pdAddrs)) {
+          tryUpdateMembers(urls);
+        }
         return;
       }
     }
@@ -458,6 +480,10 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
         "already tried all address on file, but not leader found yet.");
   }
 
+  private synchronized void tryUpdateMembers(List<URI> members) {
+    this.pdAddrs = members;
+  }
+
   public void updateTiFlashReplicaStatus() {
     ByteSequence prefix =
         ByteSequence.from(TIFLASH_TABLE_SYNC_PROGRESS_PATH, StandardCharsets.UTF_8);
@@ -513,25 +539,34 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
 
   @Override
   protected PDBlockingStub getBlockingStub() {
-    if (leaderWrapper == null) {
+    if (pdClientWrapper == null) {
       throw new GrpcException("PDClient may not be initialized");
     }
-    return leaderWrapper.getBlockingStub().withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
+    return pdClientWrapper.getBlockingStub().withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
   }
 
   @Override
   protected PDStub getAsyncStub() {
-    if (leaderWrapper == null) {
+    if (pdClientWrapper == null) {
       throw new GrpcException("PDClient may not be initialized");
     }
-    return leaderWrapper.getAsyncStub().withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
+    return pdClientWrapper.getAsyncStub().withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
   }
 
   private void initCluster() {
     GetMembersResponse resp = null;
     List<URI> pdAddrs = getConf().getPdAddrs();
     this.pdAddrs = pdAddrs;
-    this.etcdClient = Client.builder().endpoints(pdAddrs).build();
+    this.etcdClient =
+        Client.builder()
+            .endpoints(pdAddrs)
+            .executorService(
+                Executors.newCachedThreadPool(
+                    new ThreadFactoryBuilder()
+                        .setNameFormat("etcd-conn-manager-pool-%d")
+                        .setDaemon(true)
+                        .build()))
+            .build();
     this.hostMapping = new HostMapping(this.etcdClient, conf.getNetworkMappingName());
     for (URI u : pdAddrs) {
       resp = getMembers(u);
@@ -544,43 +579,65 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
     header = RequestHeader.newBuilder().setClusterId(clusterId).build();
     tsoReq = TsoRequest.newBuilder().setHeader(header).setCount(1).build();
     this.tiflashReplicaMap = new ConcurrentHashMap<>();
-    createLeaderWrapper(resp.getLeader().getClientUrls(0));
+    this.pdAddrs =
+        resp.getMembersList()
+            .stream()
+            .map(mem -> addrToUri(mem.getClientUrls(0)))
+            .collect(Collectors.toList());
+    logger.info("init cluster with address: " + this.pdAddrs);
+
+    String leaderUrlStr = resp.getLeader().getClientUrls(0);
+    leaderUrlStr = uriToAddr(addrToUri(leaderUrlStr));
+    createLeaderClientWrapper(leaderUrlStr);
     service =
         Executors.newSingleThreadScheduledExecutor(
-            new ThreadFactoryBuilder().setDaemon(true).build());
+            new ThreadFactoryBuilder()
+                .setNameFormat("PDClient-update-leader-pool-%d")
+                .setDaemon(true)
+                .build());
     service.scheduleAtFixedRate(
         () -> {
           // Wrap this with a try catch block in case schedule update fails
           try {
-            updateLeader();
+            tryUpdateLeader();
           } catch (Exception e) {
             logger.warn("Update leader failed", e);
           }
         },
-        1,
+        10,
-        1,
+        10,
-        TimeUnit.MINUTES);
+        TimeUnit.SECONDS);
     tiflashReplicaService =
         Executors.newSingleThreadScheduledExecutor(
-            new ThreadFactoryBuilder().setDaemon(true).build());
+            new ThreadFactoryBuilder()
+                .setNameFormat("PDClient-tiflash-replica-pool-%d")
+                .setDaemon(true)
+                .build());
     tiflashReplicaService.scheduleAtFixedRate(
         this::updateTiFlashReplicaStatus, 10, 10, TimeUnit.SECONDS);
   }
 
-  static class LeaderWrapper {
+  static class PDClientWrapper {
     private final String leaderInfo;
     private final PDBlockingStub blockingStub;
     private final PDStub asyncStub;
     private final long createTime;
+    private final String storeAddress;
 
-    LeaderWrapper(
+    PDClientWrapper(
-        String leaderInfo,
+        String leaderInfo, String storeAddress, ManagedChannel clientChannel, long createTime) {
-        PDGrpc.PDBlockingStub blockingStub,
+      if (!storeAddress.equals(leaderInfo)) {
-        PDGrpc.PDStub asyncStub,
+        Metadata header = new Metadata();
-        long createTime) {
+        header.put(TiConfiguration.PD_FORWARD_META_DATA_KEY, addrToUri(leaderInfo).toString());
+        this.blockingStub =
+            MetadataUtils.attachHeaders(PDGrpc.newBlockingStub(clientChannel), header);
+        this.asyncStub = MetadataUtils.attachHeaders(PDGrpc.newStub(clientChannel), header);
+      } else {
+        this.blockingStub = PDGrpc.newBlockingStub(clientChannel);
+        this.asyncStub = PDGrpc.newStub(clientChannel);
+      }
       this.leaderInfo = leaderInfo;
-      this.blockingStub = blockingStub;
+      this.storeAddress = storeAddress;
-      this.asyncStub = asyncStub;
       this.createTime = createTime;
     }
 
@@ -588,6 +645,10 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
       return leaderInfo;
     }
 
+    String getStoreAddress() {
+      return storeAddress;
+    }
+
     PDBlockingStub getBlockingStub() {
       return blockingStub;
     }

src/main/java/org/tikv/common/ReadOnlyPDClient.java

@@ -17,7 +17,6 @@ package org.tikv.common;
 
 import com.google.protobuf.ByteString;
 import java.util.List;
-import java.util.concurrent.Future;
 import org.tikv.common.meta.TiTimestamp;
 import org.tikv.common.region.TiRegion;
 import org.tikv.common.util.BackOffer;
@@ -40,8 +39,6 @@ public interface ReadOnlyPDClient {
    */
   TiRegion getRegionByKey(BackOffer backOffer, ByteString key);
 
-  Future<TiRegion> getRegionByKeyAsync(BackOffer backOffer, ByteString key);
-
   /**
    * Get Region by Region Id
    *
@@ -50,8 +47,6 @@ public interface ReadOnlyPDClient {
    */
   TiRegion getRegionByID(BackOffer backOffer, long id);
 
-  Future<TiRegion> getRegionByIDAsync(BackOffer backOffer, long id);
-
   HostMapping getHostMapping();
 
   /**
@@ -62,9 +57,7 @@ public interface ReadOnlyPDClient {
    */
   Store getStore(BackOffer backOffer, long storeId);
 
-  Future<Store> getStoreAsync(BackOffer backOffer, long storeId);
-
   List<Store> getAllStores(BackOffer backOffer);
 
-  boolean isReplicaRead();
+  TiConfiguration.ReplicaRead getReplicaRead();
 }

src/main/java/org/tikv/common/TiConfiguration.java

@@ -17,6 +17,7 @@ package org.tikv.common;
 
 import static org.tikv.common.ConfigUtils.*;
 
+import io.grpc.Metadata;
 import java.io.Serializable;
 import java.net.URI;
 import java.util.*;
@@ -24,6 +25,7 @@ import java.util.concurrent.ConcurrentHashMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.pd.PDUtils;
+import org.tikv.common.replica.ReplicaSelector;
 import org.tikv.kvproto.Kvrpcpb.CommandPri;
 import org.tikv.kvproto.Kvrpcpb.IsolationLevel;
 
@@ -31,6 +33,10 @@ public class TiConfiguration implements Serializable {
 
   private static final Logger logger = LoggerFactory.getLogger(TiConfiguration.class);
   private static final ConcurrentHashMap<String, String> settings = new ConcurrentHashMap<>();
+  public static final Metadata.Key FORWARD_META_DATA_KEY =
+      Metadata.Key.of("tikv-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
+  public static final Metadata.Key PD_FORWARD_META_DATA_KEY =
+      Metadata.Key.of("pd-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
 
   static {
     loadFromSystemProperties();
@@ -67,10 +73,12 @@ public class TiConfiguration implements Serializable {
     setIfMissing(TIKV_DB_PREFIX, DEF_DB_PREFIX);
     setIfMissing(TIKV_KV_CLIENT_CONCURRENCY, DEF_KV_CLIENT_CONCURRENCY);
     setIfMissing(TIKV_KV_MODE, TXN_KV_MODE);
-    setIfMissing(TIKV_IS_REPLICA_READ, DEF_IS_REPLICA_READ);
+    setIfMissing(TIKV_REPLICA_READ, DEF_REPLICA_READ);
     setIfMissing(TIKV_METRICS_ENABLE, DEF_METRICS_ENABLE);
     setIfMissing(TIKV_METRICS_PORT, DEF_METRICS_PORT);
     setIfMissing(TIKV_NETWORK_MAPPING_NAME, DEF_TIKV_NETWORK_MAPPING_NAME);
+    setIfMissing(TIKV_ENABLE_GRPC_FORWARD, DEF_GRPC_FORWARD_ENABLE);
+    setIfMissing(TIKV_GRPC_HEALTH_CHECK_TIMEOUT, DEF_CHECK_HEALTH_TIMEOUT);
   }
 
   public static void listAll() {
@@ -216,6 +224,17 @@ public class TiConfiguration implements Serializable {
     }
   }
 
+  private static ReplicaRead getReplicaRead(String key) {
+    String value = get(key).toUpperCase(Locale.ROOT);
+    if (FOLLOWER.equals(value)) {
+      return ReplicaRead.FOLLOWER;
+    } else if (LEADER_AND_FOLLOWER.equals(value)) {
+      return ReplicaRead.LEADER_AND_FOLLOWER;
+    } else {
+      return ReplicaRead.LEADER;
+    }
+  }
+
   private long timeout = getTimeAsMs(TIKV_GRPC_TIMEOUT);
   private long scanTimeout = getTimeAsMs(TIKV_GRPC_SCAN_TIMEOUT);
   private int maxFrameSize = getInt(TIKV_GRPC_MAX_FRAME_SIZE);
@@ -233,12 +252,16 @@ public class TiConfiguration implements Serializable {
   private boolean showRowId = getBoolean(TIKV_SHOW_ROWID);
   private String dbPrefix = get(TIKV_DB_PREFIX);
   private KVMode kvMode = getKvMode(TIKV_KV_MODE);
+  private boolean enableGrpcForward = getBoolean(TIKV_ENABLE_GRPC_FORWARD);
 
   private int kvClientConcurrency = getInt(TIKV_KV_CLIENT_CONCURRENCY);
-  private boolean isReplicaRead = getBoolean(TIKV_IS_REPLICA_READ);
+  private ReplicaRead replicaRead = getReplicaRead(TIKV_REPLICA_READ);
+  private ReplicaSelector internalReplicaSelector = getReplicaSelector(replicaRead);
+  private ReplicaSelector replicaSelector;
 
   private boolean metricsEnable = getBoolean(TIKV_METRICS_ENABLE);
   private int metricsPort = getInt(TIKV_METRICS_PORT);
+  private int grpcHealthCheckTimeout = getInt(TIKV_GRPC_HEALTH_CHECK_TIMEOUT);
 
   private final String networkMappingName = get(TIKV_NETWORK_MAPPING_NAME);
 
@@ -247,6 +270,12 @@ public class TiConfiguration implements Serializable {
     RAW
   }
 
+  public enum ReplicaRead {
+    LEADER,
+    FOLLOWER,
+    LEADER_AND_FOLLOWER
+  }
+
   public static TiConfiguration createDefault() {
     return new TiConfiguration();
   }
@@ -457,15 +486,40 @@ public class TiConfiguration implements Serializable {
     return this;
   }
 
-  public boolean isReplicaRead() {
+  public ReplicaRead getReplicaRead() {
-    return isReplicaRead;
+    return replicaRead;
   }
 
-  public TiConfiguration setReplicaRead(boolean isReplicaRead) {
+  public TiConfiguration setReplicaRead(ReplicaRead replicaRead) {
-    this.isReplicaRead = isReplicaRead;
+    this.replicaRead = replicaRead;
+    this.internalReplicaSelector = getReplicaSelector(this.replicaRead);
     return this;
   }
 
+  private ReplicaSelector getReplicaSelector(ReplicaRead replicaRead) {
+    if (TiConfiguration.ReplicaRead.LEADER.equals(replicaRead)) {
+      return ReplicaSelector.LEADER;
+    } else if (TiConfiguration.ReplicaRead.FOLLOWER.equals(replicaRead)) {
+      return ReplicaSelector.FOLLOWER;
+    } else if (TiConfiguration.ReplicaRead.LEADER_AND_FOLLOWER.equals(replicaRead)) {
+      return ReplicaSelector.LEADER_AND_FOLLOWER;
+    } else {
+      return null;
+    }
+  }
+
+  public ReplicaSelector getReplicaSelector() {
+    if (replicaSelector != null) {
+      return replicaSelector;
+    } else {
+      return internalReplicaSelector;
+    }
+  }
+
+  public void setReplicaSelector(ReplicaSelector replicaSelector) {
+    this.replicaSelector = replicaSelector;
+  }
+
   public boolean isMetricsEnable() {
     return metricsEnable;
   }
@@ -487,4 +541,12 @@ public class TiConfiguration implements Serializable {
   public String getNetworkMappingName() {
     return this.networkMappingName;
   }
+
+  public boolean getEnableGrpcForward() {
+    return this.enableGrpcForward;
+  }
+
+  public long getGrpcHealthCheckTimeout() {
+    return this.grpcHealthCheckTimeout;
+  }
 }

src/main/java/org/tikv/common/TiSession.java

@@ -43,8 +43,8 @@ import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.*;
-import org.tikv.kvproto.Metapb;
 import org.tikv.raw.RawKVClient;
 import org.tikv.txn.KVClient;
 import org.tikv.txn.TxnKVClient;
@@ -71,6 +71,7 @@ public class TiSession implements AutoCloseable {
   private volatile ExecutorService batchScanThreadPool;
   private volatile ExecutorService deleteRangeThreadPool;
   private volatile RegionManager regionManager;
+  private volatile boolean enableGrpcForward;
   private volatile RegionStoreClient.RegionStoreClientBuilder clientBuilder;
   private boolean isClosed = false;
   private HTTPServer server;
@@ -80,6 +81,7 @@ public class TiSession implements AutoCloseable {
     this.conf = conf;
     this.channelFactory = new ChannelFactory(conf.getMaxFrameSize());
     this.client = PDClient.createRaw(conf, channelFactory);
+    this.enableGrpcForward = conf.getEnableGrpcForward();
     if (conf.isMetricsEnable()) {
       try {
         this.collectorRegistry = new CollectorRegistry();
@@ -90,6 +92,7 @@ public class TiSession implements AutoCloseable {
         this.collectorRegistry.register(RetryPolicy.GRPC_SINGLE_REQUEST_LATENCY);
         this.collectorRegistry.register(RegionManager.GET_REGION_BY_KEY_REQUEST_LATENCY);
         this.collectorRegistry.register(PDClient.PD_GET_REGION_BY_KEY_REQUEST_LATENCY);
+        this.enableGrpcForward = conf.getEnableGrpcForward();
         this.server =
             new HTTPServer(
                 new InetSocketAddress(conf.getMetricsPort()), this.collectorRegistry, true);
@@ -194,12 +197,17 @@ public class TiSession implements AutoCloseable {
     return res;
   }
 
-  public synchronized RegionManager getRegionManager() {
+  public RegionManager getRegionManager() {
     RegionManager res = regionManager;
     if (res == null) {
       synchronized (this) {
         if (regionManager == null) {
-          regionManager = new RegionManager(getPDClient(), this.cacheInvalidateCallback);
+          regionManager =
+              new RegionManager(
+                  getPDClient(),
+                  this.cacheInvalidateCallback,
+                  this.channelFactory,
+                  this.enableGrpcForward);
         }
         res = regionManager;
       }
@@ -415,10 +423,10 @@ public class TiSession implements AutoCloseable {
         groupKeysByRegion(regionManager, splitKeys, backOffer);
     for (Map.Entry<TiRegion, List<ByteString>> entry : groupKeys.entrySet()) {
 
-      Pair<TiRegion, Metapb.Store> pair =
+      Pair<TiRegion, TiStore> pair =
           getRegionManager().getRegionStorePairByKey(entry.getKey().getStartKey());
       TiRegion region = pair.first;
-      Metapb.Store store = pair.second;
+      TiStore store = pair.second;
       List<ByteString> splits =
           entry
               .getValue()
@@ -466,7 +474,9 @@ public class TiSession implements AutoCloseable {
     synchronized (sessionCachedMap) {
       sessionCachedMap.remove(conf.getPdAddrsString());
     }
-
+    if (regionManager != null) {
+      regionManager.close();
+    }
     if (tableScanThreadPool != null) {
       tableScanThreadPool.shutdownNow();
     }

src/main/java/org/tikv/common/allocator/RowIDAllocator.java

@@ -1,258 +0,0 @@
-/*
- * Copyright 2019 PingCAP, Inc.
- *
- * 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,
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.tikv.common.allocator;
-
-import com.google.common.primitives.UnsignedLongs;
-import com.google.protobuf.ByteString;
-import java.io.Serializable;
-import java.util.Arrays;
-import java.util.function.Function;
-import org.tikv.common.Snapshot;
-import org.tikv.common.TiSession;
-import org.tikv.common.codec.CodecDataInput;
-import org.tikv.common.codec.CodecDataOutput;
-import org.tikv.common.codec.MetaCodec;
-import org.tikv.common.exception.AllocateRowIDOverflowException;
-import org.tikv.common.exception.TiBatchWriteException;
-import org.tikv.common.meta.TiTableInfo;
-import org.tikv.common.util.BackOffer;
-import org.tikv.common.util.ConcreteBackOffer;
-import org.tikv.txn.TwoPhaseCommitter;
-
-/**
- * RowIDAllocator read current start from TiKV and write back 'start+step' back to TiKV. It designs
- * to allocate all id for data to be written at once, hence it does not need run inside a txn.
- *
- * <p>(start, end] is allocated
- */
-public final class RowIDAllocator implements Serializable {
-  private final long maxShardRowIDBits;
-  private final long dbId;
-  private final TiSession session;
-  private final long step;
-  private long end;
-
-  private RowIDAllocator(long maxShardRowIDBits, long dbId, long step, TiSession session) {
-    this.maxShardRowIDBits = maxShardRowIDBits;
-    this.dbId = dbId;
-    this.step = step;
-    this.session = session;
-  }
-
-  /**
-   * @param index should >= 1
-   * @return
-   */
-  public long getShardRowId(long index) {
-    return getShardRowId(maxShardRowIDBits, index, index + getStart());
-  }
-
-  static long getShardRowId(long maxShardRowIDBits, long partitionIndex, long rowID) {
-    if (maxShardRowIDBits <= 0 || maxShardRowIDBits >= 16) {
-      return rowID;
-    }
-
-    // assert rowID < Math.pow(2, 64 - maxShardRowIDBits)
-
-    long partition = partitionIndex & ((1L << maxShardRowIDBits) - 1);
-    return rowID | (partition << (64 - maxShardRowIDBits - 1));
-  }
-
-  public static RowIDAllocator create(
-      long dbId, TiTableInfo table, TiSession session, boolean unsigned, long step) {
-    RowIDAllocator allocator =
-        new RowIDAllocator(table.getMaxShardRowIDBits(), dbId, step, session);
-    if (unsigned) {
-      allocator.initUnsigned(session.createSnapshot(), table.getId(), table.getMaxShardRowIDBits());
-    } else {
-      allocator.initSigned(session.createSnapshot(), table.getId(), table.getMaxShardRowIDBits());
-    }
-
-    return allocator;
-  }
-
-  public long getStart() {
-    return end - step;
-  }
-
-  public long getEnd() {
-    return end;
-  }
-
-  // set key value pair to tikv via two phase committer protocol.
-  private void set(ByteString key, byte[] value) {
-    TwoPhaseCommitter twoPhaseCommitter =
-        new TwoPhaseCommitter(session, session.getTimestamp().getVersion());
-
-    twoPhaseCommitter.prewritePrimaryKey(
-        ConcreteBackOffer.newCustomBackOff(BackOffer.PREWRITE_MAX_BACKOFF),
-        key.toByteArray(),
-        value);
-
-    twoPhaseCommitter.commitPrimaryKey(
-        ConcreteBackOffer.newCustomBackOff(BackOffer.BATCH_COMMIT_BACKOFF),
-        key.toByteArray(),
-        session.getTimestamp().getVersion());
-
-    try {
-      twoPhaseCommitter.close();
-    } catch (Throwable ignored) {
-    }
-  }
-
-  private void updateMeta(ByteString key, byte[] oldVal, Snapshot snapshot) {
-    // 1. encode hash meta key
-    // 2. load meta via hash meta key from TiKV
-    // 3. update meta's filed count and set it back to TiKV
-    CodecDataOutput cdo = new CodecDataOutput();
-    ByteString metaKey = MetaCodec.encodeHashMetaKey(cdo, key.toByteArray());
-    long fieldCount;
-    ByteString metaVal = snapshot.get(metaKey);
-
-    // decode long from bytes
-    // big endian the 8 bytes
-    fieldCount = new CodecDataInput(metaVal.toByteArray()).readLong();
-
-    // update meta field count only oldVal is null
-    if (oldVal == null || oldVal.length == 0) {
-      fieldCount++;
-      cdo.reset();
-      cdo.writeLong(fieldCount);
-
-      set(metaKey, cdo.toBytes());
-    }
-  }
-
-  private long updateHash(
-      ByteString key,
-      ByteString field,
-      Function<byte[], byte[]> calculateNewVal,
-      Snapshot snapshot) {
-    // 1. encode hash data key
-    // 2. get value in byte from get operation
-    // 3. calculate new value via calculateNewVal
-    // 4. check old value equals to new value or not
-    // 5. set the new value back to TiKV via 2pc
-    // 6. encode a hash meta key
-    // 7. update a hash meta field count if needed
-
-    CodecDataOutput cdo = new CodecDataOutput();
-    MetaCodec.encodeHashDataKey(cdo, key.toByteArray(), field.toByteArray());
-    ByteString dataKey = cdo.toByteString();
-    byte[] oldVal = snapshot.get(dataKey.toByteArray());
-
-    byte[] newVal = calculateNewVal.apply(oldVal);
-    if (Arrays.equals(newVal, oldVal)) {
-      // not need to update
-      return 0L;
-    }
-
-    set(dataKey, newVal);
-    updateMeta(key, oldVal, snapshot);
-    return Long.parseLong(new String(newVal));
-  }
-
-  private static boolean isDBExisted(long dbId, Snapshot snapshot) {
-    ByteString dbKey = MetaCodec.encodeDatabaseID(dbId);
-    ByteString json = MetaCodec.hashGet(MetaCodec.KEY_DBs, dbKey, snapshot);
-    return json != null && !json.isEmpty();
-  }
-
-  private static boolean isTableExisted(long dbId, long tableId, Snapshot snapshot) {
-    ByteString dbKey = MetaCodec.encodeDatabaseID(dbId);
-    ByteString tableKey = MetaCodec.tableKey(tableId);
-    return !MetaCodec.hashGet(dbKey, tableKey, snapshot).isEmpty();
-  }
-
-  public static boolean shardRowBitsOverflow(
-      long base, long step, long shardRowBits, boolean reservedSignBit) {
-    long signBit = reservedSignBit ? 1 : 0;
-    long mask = ((1L << shardRowBits) - 1) << (64 - shardRowBits - signBit);
-    if (reservedSignBit) {
-      return ((base + step) & mask) > 0;
-    } else {
-      return Long.compareUnsigned((base + step) & mask, 0) > 0;
-    }
-  }
-
-  /**
-   * read current row id from TiKV and write the calculated value back to TiKV. The calculation rule
-   * is start(read from TiKV) + step.
-   */
-  public long udpateAllocateId(
-      long dbId, long tableId, long step, Snapshot snapshot, long shard, boolean hasSignedBit) {
-    if (isDBExisted(dbId, snapshot) && isTableExisted(dbId, tableId, snapshot)) {
-      return updateHash(
-          MetaCodec.encodeDatabaseID(dbId),
-          MetaCodec.autoTableIDKey(tableId),
-          (oldVal) -> {
-            long base = 0;
-            if (oldVal != null && oldVal.length != 0) {
-              base = Long.parseLong(new String(oldVal));
-            }
-            if (shard >= 1 && shardRowBitsOverflow(base, step, shard, hasSignedBit)) {
-              throw new AllocateRowIDOverflowException(base, step, shard);
-            }
-            base += step;
-            return String.valueOf(base).getBytes();
-          },
-          snapshot);
-    }
-
-    throw new IllegalArgumentException("table or database is not existed");
-  }
-
-  /** read current row id from TiKV according to database id and table id. */
-  public static long getAllocateId(long dbId, long tableId, Snapshot snapshot) {
-    if (isDBExisted(dbId, snapshot) && isTableExisted(dbId, tableId, snapshot)) {
-      ByteString dbKey = MetaCodec.encodeDatabaseID(dbId);
-      ByteString tblKey = MetaCodec.autoTableIDKey(tableId);
-      ByteString val = MetaCodec.hashGet(dbKey, tblKey, snapshot);
-      if (val.isEmpty()) return 0L;
-      return Long.parseLong(val.toStringUtf8());
-    }
-
-    throw new IllegalArgumentException("table or database is not existed");
-  }
-
-  private void initSigned(Snapshot snapshot, long tableId, long shard) {
-    // get new start from TiKV, and calculate new end and set it back to TiKV.
-    long newStart = getAllocateId(dbId, tableId, snapshot);
-    long tmpStep = Math.min(Long.MAX_VALUE - newStart, step);
-    if (tmpStep != step) {
-      throw new TiBatchWriteException("cannot allocate ids for this write");
-    }
-    if (newStart == Long.MAX_VALUE) {
-      throw new TiBatchWriteException("cannot allocate more ids since it ");
-    }
-    end = udpateAllocateId(dbId, tableId, tmpStep, snapshot, shard, true);
-  }
-
-  private void initUnsigned(Snapshot snapshot, long tableId, long shard) {
-    // get new start from TiKV, and calculate new end and set it back to TiKV.
-    long newStart = getAllocateId(dbId, tableId, snapshot);
-    // for unsigned long, -1L is max value.
-    long tmpStep = UnsignedLongs.min(-1L - newStart, step);
-    if (tmpStep != step) {
-      throw new TiBatchWriteException("cannot allocate ids for this write");
-    }
-    // when compare unsigned long, the min value is largest value.
-    if (UnsignedLongs.compare(newStart, -1L) == 0) {
-      throw new TiBatchWriteException(
-          "cannot allocate more ids since the start reaches " + "unsigned long's max value ");
-    }
-    end = udpateAllocateId(dbId, tableId, tmpStep, snapshot, shard, false);
-  }
-}

src/main/java/org/tikv/common/exception/RegionException.java

@@ -24,7 +24,7 @@ public class RegionException extends TiKVException {
   private final Error regionErr;
 
   public RegionException(Error regionErr) {
-    super("Region Exception occurred" + regionErr.getMessage());
+    super("Region Exception occurred " + regionErr.getMessage());
     this.regionErr = regionErr;
   }
 

src/main/java/org/tikv/common/operation/KVErrorHandler.java

@@ -123,9 +123,7 @@ public class KVErrorHandler<RespT> implements ErrorHandler<RespT> {
   @Override
   public boolean handleResponseError(BackOffer backOffer, RespT resp) {
     if (resp == null) {
-      String msg =
+      String msg = String.format("Request Failed with unknown reason for [%s]", recv.getRegion());
-          String.format(
-              "Request Failed with unknown reason for region region [%s]", recv.getRegion());
       logger.warn(msg);
       return handleRequestError(backOffer, new GrpcException(msg));
     }
@@ -158,9 +156,7 @@ public class KVErrorHandler<RespT> implements ErrorHandler<RespT> {
           // onNotLeader is only needed when updateLeader succeeds, thus switch
           // to a new store address.
           TiRegion newRegion = this.regionManager.updateLeader(recv.getRegion(), newStoreId);
-          retry =
+          retry = newRegion != null && recv.onNotLeader(newRegion);
-              newRegion != null
-                  && recv.onNotLeader(this.regionManager.getStoreById(newStoreId), newRegion);
 
           backOffFuncType = BackOffFunction.BackOffFuncType.BoUpdateLeader;
         } else {

src/main/java/org/tikv/common/operation/PDErrorHandler.java

@@ -48,7 +48,9 @@ public class PDErrorHandler<RespT> implements ErrorHandler<RespT> {
   @Override
   public boolean handleResponseError(BackOffer backOffer, RespT resp) {
     if (resp == null) {
-      return false;
+      String msg = String.format("PD Request Failed with unknown reason");
+      logger.warn(msg);
+      return handleRequestError(backOffer, new GrpcException(msg));
     }
     PDError error = getError.apply(resp);
     if (error != null) {
@@ -56,7 +58,7 @@ public class PDErrorHandler<RespT> implements ErrorHandler<RespT> {
         case PD_ERROR:
           backOffer.doBackOff(
               BackOffFunction.BackOffFuncType.BoPDRPC, new GrpcException(error.toString()));
-          client.updateLeader();
+          client.updateLeaderOrforwardFollower();
           return true;
         case REGION_PEER_NOT_ELECTED:
           logger.debug(error.getMessage());
@@ -73,6 +75,7 @@ public class PDErrorHandler<RespT> implements ErrorHandler<RespT> {
   @Override
   public boolean handleRequestError(BackOffer backOffer, Exception e) {
     backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoPDRPC, e);
+    client.updateLeaderOrforwardFollower();
     return true;
   }
 }

src/main/java/org/tikv/common/operation/iterator/ConcreteScanIterator.java

@@ -27,11 +27,11 @@ import org.tikv.common.key.Key;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.Pair;
 import org.tikv.kvproto.Kvrpcpb;
-import org.tikv.kvproto.Metapb;
 
 public class ConcreteScanIterator extends ScanIterator {
   private final long version;
@@ -82,10 +82,10 @@ public class ConcreteScanIterator extends ScanIterator {
 
   private ByteString resolveCurrentLock(Kvrpcpb.KvPair current) {
     logger.warn(String.format("resolve current key error %s", current.getError().toString()));
-    Pair<TiRegion, Metapb.Store> pair =
+    Pair<TiRegion, TiStore> pair =
         builder.getRegionManager().getRegionStorePairByKey(current.getKey());
     TiRegion region = pair.first;
-    Metapb.Store store = pair.second;
+    TiStore store = pair.second;
     BackOffer backOffer = ConcreteBackOffer.newGetBackOff();
     try (RegionStoreClient client = builder.build(region, store)) {
       return client.get(backOffer, current.getKey(), version);

src/main/java/org/tikv/common/operation/iterator/CoprocessorIterator.java

@@ -27,7 +27,11 @@ import java.util.List;
 import org.tikv.common.TiSession;
 import org.tikv.common.codec.Codec.IntegerCodec;
 import org.tikv.common.codec.CodecDataInput;
-import org.tikv.common.columnar.*;
+import org.tikv.common.columnar.BatchedTiChunkColumnVector;
+import org.tikv.common.columnar.TiChunk;
+import org.tikv.common.columnar.TiChunkColumnVector;
+import org.tikv.common.columnar.TiColumnVector;
+import org.tikv.common.columnar.TiRowColumnVector;
 import org.tikv.common.columnar.datatypes.CHType;
 import org.tikv.common.meta.TiDAGRequest;
 import org.tikv.common.operation.SchemaInfer;

src/main/java/org/tikv/common/operation/iterator/DAGIterator.java

@@ -32,12 +32,12 @@ import org.tikv.common.meta.TiDAGRequest.PushDownType;
 import org.tikv.common.operation.SchemaInfer;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.region.TiStoreType;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.RangeSplitter;
 import org.tikv.kvproto.Coprocessor;
-import org.tikv.kvproto.Metapb;
 
 public abstract class DAGIterator<T>
     extends org.tikv.common.operation.iterator.CoprocessorIterator<T> {
@@ -204,7 +204,7 @@ public abstract class DAGIterator<T>
       }
       List<Coprocessor.KeyRange> ranges = task.getRanges();
       TiRegion region = task.getRegion();
-      Metapb.Store store = task.getStore();
+      TiStore store = task.getStore();
 
       try {
         RegionStoreClient client =
@@ -246,7 +246,7 @@ public abstract class DAGIterator<T>
   private Iterator<SelectResponse> processByStreaming(RangeSplitter.RegionTask regionTask) {
     List<Coprocessor.KeyRange> ranges = regionTask.getRanges();
     TiRegion region = regionTask.getRegion();
-    Metapb.Store store = regionTask.getStore();
+    TiStore store = regionTask.getStore();
 
     RegionStoreClient client;
     try {

src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java

@@ -78,26 +78,22 @@ public class RawScanIterator extends ScanIterator {
       endOfScan = true;
       return false;
     }
+    // continue when cache is empty but not null
+    while (currentCache != null && currentCache.isEmpty()) {
+      if (cacheLoadFails()) {
+        return false;
+      }
+    }
     return notEndOfScan();
   }
 
   private Kvrpcpb.KvPair getCurrent() {
-    if (isCacheDrained()) {
-      return null;
-    }
     --limit;
     return currentCache.get(index++);
   }
 
   @Override
   public Kvrpcpb.KvPair next() {
-    Kvrpcpb.KvPair kv;
+    return getCurrent();
-    // continue when cache is empty but not null
-    for (kv = getCurrent(); currentCache != null && kv == null; kv = getCurrent()) {
-      if (cacheLoadFails()) {
-        return null;
-      }
-    }
-    return kv;
   }
 }

src/main/java/org/tikv/common/operation/iterator/ScanIterator.java

@@ -50,11 +50,8 @@ public abstract class ScanIterator implements Iterator<Kvrpcpb.KvPair> {
       int limit,
       boolean keyOnly) {
     this.startKey = requireNonNull(startKey, "start key is null");
-    if (startKey.isEmpty()) {
-      throw new IllegalArgumentException("start key cannot be empty");
-    }
     this.endKey = Key.toRawKey(requireNonNull(endKey, "end key is null"));
-    this.hasEndKey = !endKey.equals(ByteString.EMPTY);
+    this.hasEndKey = !endKey.isEmpty();
     this.limit = limit;
     this.keyOnly = keyOnly;
     this.conf = conf;
@@ -74,7 +71,7 @@ public abstract class ScanIterator implements Iterator<Kvrpcpb.KvPair> {
     if (endOfScan || processingLastBatch) {
       return true;
     }
-    if (startKey == null || startKey.isEmpty()) {
+    if (startKey == null) {
       return true;
     }
     try {
@@ -107,7 +104,8 @@ public abstract class ScanIterator implements Iterator<Kvrpcpb.KvPair> {
         startKey = lastKey.next().toByteString();
       }
       // notify last batch if lastKey is greater than or equal to endKey
-      if (hasEndKey && lastKey.compareTo(endKey) >= 0) {
+      // if startKey is empty, it indicates +∞
+      if (hasEndKey && lastKey.compareTo(endKey) >= 0 || startKey.isEmpty()) {
         processingLastBatch = true;
         startKey = null;
       }

src/main/java/org/tikv/common/policy/RetryPolicy.java

@@ -17,6 +17,7 @@ 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;
@@ -32,6 +33,12 @@ public abstract class RetryPolicy<RespT> {
           .help("grpc request latency.")
           .labelNames("type")
           .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")
+          .register();
 
   // handles PD and TiKV's error.
   private ErrorHandler<RespT> handler;
@@ -70,6 +77,7 @@ public abstract class RetryPolicy<RespT> {
         // Handle request call error
         boolean retry = handler.handleRequestError(backOffer, e);
         if (retry) {
+          GRPC_REQUEST_RETRY_NUM.labels(methodName).inc();
           continue;
         }
       }
@@ -78,7 +86,7 @@ public abstract class RetryPolicy<RespT> {
       if (handler != null) {
         boolean retry = handler.handleResponseError(backOffer, result);
         if (retry) {
-          // add retry counter
+          GRPC_REQUEST_RETRY_NUM.labels(methodName).inc();
           continue;
         }
       }

src/main/java/org/tikv/common/region/AbstractRegionStoreClient.java

@@ -21,6 +21,12 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.health.v1.HealthCheckRequest;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
+import io.grpc.stub.MetadataUtils;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import org.tikv.common.AbstractGRPCClient;
 import org.tikv.common.TiConfiguration;
@@ -35,10 +41,12 @@ public abstract class AbstractRegionStoreClient
 
   protected final RegionManager regionManager;
   protected TiRegion region;
+  protected TiStore targetStore;
 
   protected AbstractRegionStoreClient(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       ChannelFactory channelFactory,
       TikvGrpc.TikvBlockingStub blockingStub,
       TikvGrpc.TikvStub asyncStub,
@@ -49,6 +57,7 @@ public abstract class AbstractRegionStoreClient
     checkArgument(region.getLeader() != null, "Leader Peer is null");
     this.region = region;
     this.regionManager = regionManager;
+    this.targetStore = store;
   }
 
   public TiRegion getRegion() {
@@ -71,13 +80,13 @@ public abstract class AbstractRegionStoreClient
   /**
    * onNotLeader deals with NotLeaderError and returns whether re-splitting key range is needed
    *
-   * @param newStore the new store presented by NotLeader Error
+   * @param newRegion the new region presented by NotLeader Error
    * @return false when re-split is needed.
    */
   @Override
-  public boolean onNotLeader(Metapb.Store newStore, TiRegion newRegion) {
+  public boolean onNotLeader(TiRegion newRegion) {
     if (logger.isDebugEnabled()) {
-      logger.debug(region + ", new leader = " + newStore.getId());
+      logger.debug(region + ", new leader = " + newRegion.getLeader().getStoreId());
     }
     // When switch leader fails or the region changed its region epoch,
     // it would be necessary to re-split task's key range for new region.
@@ -85,7 +94,8 @@ public abstract class AbstractRegionStoreClient
       return false;
     }
     region = newRegion;
-    String addressStr = regionManager.getStoreById(region.getLeader().getStoreId()).getAddress();
+    targetStore = regionManager.getStoreById(region.getLeader().getStoreId());
+    String addressStr = targetStore.getStore().getAddress();
     ManagedChannel channel =
         channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
     blockingStub = TikvGrpc.newBlockingStub(channel);
@@ -94,20 +104,84 @@ public abstract class AbstractRegionStoreClient
   }
 
   @Override
-  public void onStoreNotMatch(Metapb.Store store) {
+  public boolean onStoreUnreachable() {
-    String addressStr = store.getAddress();
+    if (!conf.getEnableGrpcForward()) {
+      return false;
+    }
+    if (region.getProxyStore() == null) {
+      if (!targetStore.isUnreachable()) {
+        if (checkHealth(targetStore)) {
+          return true;
+        } else {
+          if (targetStore.markUnreachable()) {
+            this.regionManager.scheduleHealthCheckJob(targetStore);
+          }
+        }
+      }
+    }
+    TiRegion proxyRegion = switchProxyStore();
+    if (proxyRegion == null) {
+      return false;
+    }
+    regionManager.updateRegion(region, proxyRegion);
+    region = proxyRegion;
+    String addressStr = region.getProxyStore().getStore().getAddress();
     ManagedChannel channel =
         channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
-    blockingStub = TikvGrpc.newBlockingStub(channel);
+    Metadata header = new Metadata();
-    asyncStub = TikvGrpc.newStub(channel);
+    header.put(TiConfiguration.FORWARD_META_DATA_KEY, targetStore.getStore().getAddress());
-    if (region.getLeader().getStoreId() != store.getId()) {
+    blockingStub = MetadataUtils.attachHeaders(TikvGrpc.newBlockingStub(channel), header);
-      logger.warn(
+    asyncStub = MetadataUtils.attachHeaders(TikvGrpc.newStub(channel), header);
-          "store_not_match may occur? "
+    return true;
-              + region
+  }
-              + ", original store = "
+
-              + store.getId()
+  private boolean checkHealth(TiStore store) {
-              + " address = "
+    if (store.getStore() == null) {
-              + addressStr);
+      return false;
     }
+    String addressStr = store.getStore().getAddress();
+    ManagedChannel channel =
+        channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
+    HealthGrpc.HealthBlockingStub stub =
+        HealthGrpc.newBlockingStub(channel)
+            .withDeadlineAfter(conf.getGrpcHealthCheckTimeout(), TimeUnit.MILLISECONDS);
+    HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
+    try {
+      HealthCheckResponse resp = stub.check(req);
+      if (resp.getStatus() != HealthCheckResponse.ServingStatus.SERVING) {
+        return false;
+      }
+    } catch (Exception e) {
+      return false;
+    }
+    return true;
+  }
+
+  private TiRegion switchProxyStore() {
+    boolean hasVisitedStore = false;
+    List<Metapb.Peer> peers = region.getFollowerList();
+    for (int i = 0; i < peers.size() * 2; i++) {
+      int idx = i % peers.size();
+      Metapb.Peer peer = peers.get(idx);
+      if (peer.getStoreId() != region.getLeader().getStoreId()) {
+        if (region.getProxyStore() == null) {
+          TiStore store = regionManager.getStoreById(peer.getStoreId());
+          if (checkHealth(store)) {
+            return region.switchProxyStore(store);
+          }
+        } else {
+          TiStore proxyStore = region.getProxyStore();
+          if (peer.getStoreId() == proxyStore.getStore().getId()) {
+            hasVisitedStore = true;
+          } else if (hasVisitedStore) {
+            proxyStore = regionManager.getStoreById(peer.getStoreId());
+            if (!proxyStore.isUnreachable() && checkHealth(proxyStore)) {
+              return region.switchProxyStore(proxyStore);
+            }
+          }
+        }
+      }
+    }
+    return null;
   }
 }

src/main/java/org/tikv/common/region/RegionErrorReceiver.java

@@ -17,12 +17,11 @@
 
 package org.tikv.common.region;
 
-import org.tikv.kvproto.Metapb.Store;
-
 public interface RegionErrorReceiver {
-  boolean onNotLeader(Store store, TiRegion region);
+  boolean onNotLeader(TiRegion region);
 
-  void onStoreNotMatch(Store store);
+  /// return whether we need to retry this request.
+  boolean onStoreUnreachable();
 
   TiRegion getRegion();
 }

src/main/java/org/tikv/common/region/RegionManager.java

@@ -28,6 +28,9 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -37,11 +40,11 @@ import org.tikv.common.exception.GrpcException;
 import org.tikv.common.exception.TiClientInternalException;
 import org.tikv.common.key.Key;
 import org.tikv.common.util.BackOffer;
+import org.tikv.common.util.ChannelFactory;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.Pair;
 import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.Metapb.Peer;
-import org.tikv.kvproto.Metapb.Store;
 import org.tikv.kvproto.Metapb.StoreState;
 
 @SuppressWarnings("UnstableApiUsage")
@@ -50,7 +53,8 @@ public class RegionManager {
   // TODO: the region cache logic need rewrite.
   // https://github.com/pingcap/tispark/issues/1170
   private final RegionCache cache;
-  private final boolean isReplicaRead;
+  private final ScheduledExecutorService executor;
+  private final UnreachableStoreChecker storeChecker;
 
   private final Function<CacheInvalidateEvent, Void> cacheInvalidateCallback;
 
@@ -65,14 +69,40 @@ public class RegionManager {
   public RegionManager(
       ReadOnlyPDClient pdClient, Function<CacheInvalidateEvent, Void> cacheInvalidateCallback) {
     this.cache = new RegionCache(pdClient);
-    this.isReplicaRead = pdClient.isReplicaRead();
     this.cacheInvalidateCallback = cacheInvalidateCallback;
+    this.executor = null;
+    this.storeChecker = null;
+  }
+
+  public RegionManager(
+      ReadOnlyPDClient pdClient,
+      Function<CacheInvalidateEvent, Void> cacheInvalidateCallback,
+      ChannelFactory channelFactory,
+      boolean enableGrpcForward) {
+    this.cache = new RegionCache(pdClient);
+    this.cacheInvalidateCallback = cacheInvalidateCallback;
+    if (enableGrpcForward) {
+      UnreachableStoreChecker storeChecker = new UnreachableStoreChecker(channelFactory, pdClient);
+      this.storeChecker = storeChecker;
+      this.executor = Executors.newScheduledThreadPool(1);
+      this.executor.scheduleAtFixedRate(storeChecker, 10, 10, TimeUnit.SECONDS);
+    } else {
+      this.storeChecker = null;
+      this.executor = null;
+    }
   }
 
   public RegionManager(ReadOnlyPDClient pdClient) {
     this.cache = new RegionCache(pdClient);
-    this.isReplicaRead = pdClient.isReplicaRead();
     this.cacheInvalidateCallback = null;
+    this.storeChecker = null;
+    this.executor = null;
+  }
+
+  public synchronized void close() {
+    if (this.executor != null) {
+      this.executor.shutdownNow();
+    }
   }
 
   public Function<CacheInvalidateEvent, Void> getCacheInvalidateCallback() {
@@ -102,19 +132,19 @@ public class RegionManager {
     return cache.getRegionById(ConcreteBackOffer.newGetBackOff(), regionId);
   }
 
-  public Pair<TiRegion, Store> getRegionStorePairByKey(ByteString key, BackOffer backOffer) {
+  public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key, BackOffer backOffer) {
     return getRegionStorePairByKey(key, TiStoreType.TiKV, backOffer);
   }
 
-  public Pair<TiRegion, Store> getRegionStorePairByKey(ByteString key) {
+  public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key) {
     return getRegionStorePairByKey(key, TiStoreType.TiKV);
   }
 
-  public Pair<TiRegion, Store> getRegionStorePairByKey(ByteString key, TiStoreType storeType) {
+  public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key, TiStoreType storeType) {
     return getRegionStorePairByKey(key, storeType, ConcreteBackOffer.newGetBackOff());
   }
 
-  public Pair<TiRegion, Store> getRegionStorePairByKey(
+  public Pair<TiRegion, TiStore> getRegionStorePairByKey(
       ByteString key, TiStoreType storeType, BackOffer backOffer) {
     TiRegion region = cache.getRegionByKey(key, backOffer);
     if (region == null) {
@@ -124,20 +154,15 @@ public class RegionManager {
       throw new TiClientInternalException("Region invalid: " + region.toString());
     }
 
-    Store store = null;
+    TiStore store = null;
     if (storeType == TiStoreType.TiKV) {
-      if (isReplicaRead) {
+      Peer peer = region.getCurrentReplica();
-        Peer peer = region.getCurrentFollower();
+      store = cache.getStoreById(peer.getStoreId(), backOffer);
-        store = cache.getStoreById(peer.getStoreId(), backOffer);
-      } else {
-        Peer leader = region.getLeader();
-        store = cache.getStoreById(leader.getStoreId(), backOffer);
-      }
     } else {
       outerLoop:
       for (Peer peer : region.getLearnerList()) {
-        Store s = getStoreById(peer.getStoreId(), backOffer);
+        TiStore s = getStoreById(peer.getStoreId(), backOffer);
-        for (Metapb.StoreLabel label : s.getLabelsList()) {
+        for (Metapb.StoreLabel label : s.getStore().getLabelsList()) {
           if (label.getKey().equals(storeType.getLabelKey())
               && label.getValue().equals(storeType.getLabelValue())) {
             store = s;
@@ -159,11 +184,11 @@ public class RegionManager {
     return Pair.create(region, store);
   }
 
-  public Store getStoreById(long id) {
+  public TiStore getStoreById(long id) {
     return getStoreById(id, ConcreteBackOffer.newGetBackOff());
   }
 
-  public Store getStoreById(long id, BackOffer backOffer) {
+  public TiStore getStoreById(long id, BackOffer backOffer) {
     return cache.getStoreById(id, backOffer);
   }
 
@@ -172,23 +197,25 @@ public class RegionManager {
   }
 
   public synchronized TiRegion updateLeader(TiRegion region, long storeId) {
-    TiRegion r = cache.getRegionFromCache(region.getId());
+    TiRegion newRegion = region.switchPeer(storeId);
-    if (r != null) {
+    if (cache.updateRegion(region, newRegion)) {
-      if (r.getLeader().getStoreId() == storeId) {
+      return newRegion;
-        return r;
-      }
-      TiRegion newRegion = r.switchPeer(storeId);
-      if (newRegion != null) {
-        cache.putRegion(newRegion);
-        return newRegion;
-      }
-      // failed to switch leader, possibly region is outdated, we need to drop region cache from
-      // regionCache
-      logger.warn("Cannot find peer when updating leader (" + region.getId() + "," + storeId + ")");
     }
+    // failed to switch leader, possibly region is outdated, we need to drop region cache from
+    // regionCache
+    logger.warn("Cannot find peer when updating leader (" + region.getId() + "," + storeId + ")");
     return null;
   }
 
+  public synchronized boolean updateRegion(TiRegion oldRegion, TiRegion region) {
+    return cache.updateRegion(oldRegion, region);
+  }
+
+  /** Clears all cache when some unexpected error occurs. */
+  public void clearRegionCache() {
+    cache.clearAll();
+  }
+
   /**
    * Clears all cache when a TiKV server does not respond
    *
@@ -199,8 +226,10 @@ public class RegionManager {
   }
 
   private void onRequestFail(TiRegion region, long storeId) {
-    cache.invalidateRegion(region);
+    if (this.storeChecker != null) {
-    cache.invalidateAllRegionForStore(storeId);
+      cache.invalidateRegion(region);
+      cache.invalidateAllRegionForStore(storeId);
+    }
   }
 
   public void invalidateStore(long storeId) {
@@ -211,9 +240,13 @@ public class RegionManager {
     cache.invalidateRegion(region);
   }
 
+  public void scheduleHealthCheckJob(TiStore store) {
+    this.storeChecker.scheduleStoreHealthCheck(store);
+  }
+
   public static class RegionCache {
     private final Map<Long, TiRegion> regionCache;
-    private final Map<Long, Store> storeCache;
+    private final Map<Long, TiStore> storeCache;
     private final RangeMap<Key, Long> keyToRegionIdCache;
     private final ReadOnlyPDClient pdClient;
 
@@ -303,6 +336,26 @@ public class RegionManager {
       }
     }
 
+    public synchronized boolean updateRegion(TiRegion expected, TiRegion region) {
+      try {
+        if (logger.isDebugEnabled()) {
+          logger.debug(String.format("invalidateRegion ID[%s]", region.getId()));
+        }
+        TiRegion oldRegion = regionCache.get(region.getId());
+        if (expected != oldRegion) {
+          return false;
+        } else {
+          if (oldRegion != null) {
+            keyToRegionIdCache.remove(makeRange(oldRegion.getStartKey(), oldRegion.getEndKey()));
+          }
+          putRegion(region);
+          return true;
+        }
+      } catch (Exception ignore) {
+        return false;
+      }
+    }
+
     public synchronized void invalidateAllRegionForStore(long storeId) {
       List<TiRegion> regionToRemove = new ArrayList<>();
       for (TiRegion r : regionCache.values()) {
@@ -316,22 +369,25 @@ public class RegionManager {
 
       // remove region
       for (TiRegion r : regionToRemove) {
-        regionCache.remove(r.getId());
         keyToRegionIdCache.remove(makeRange(r.getStartKey(), r.getEndKey()));
+        regionCache.remove(r.getId());
       }
     }
 
     public synchronized void invalidateStore(long storeId) {
-      storeCache.remove(storeId);
+      TiStore store = storeCache.remove(storeId);
+      if (store != null) {
+        store.markReachable();
+      }
     }
 
-    public synchronized Store getStoreById(long id, BackOffer backOffer) {
+    public synchronized TiStore getStoreById(long id, BackOffer backOffer) {
       try {
-        Store store = storeCache.get(id);
+        TiStore store = storeCache.get(id);
         if (store == null) {
-          store = pdClient.getStore(backOffer, id);
+          store = new TiStore(pdClient.getStore(backOffer, id));
         }
-        if (store.getState().equals(StoreState.Tombstone)) {
+        if (store.getStore().getState().equals(StoreState.Tombstone)) {
           return null;
         }
         storeCache.put(id, store);
@@ -340,5 +396,10 @@ public class RegionManager {
         throw new GrpcException(e);
       }
     }
+
+    public synchronized void clearAll() {
+      keyToRegionIdCache.clear();
+      regionCache.clear();
+    }
   }
 }

src/main/java/org/tikv/common/region/RegionStoreClient.java

@@ -26,6 +26,8 @@ import com.google.protobuf.InvalidProtocolBufferException;
 import com.pingcap.tidb.tipb.DAGRequest;
 import com.pingcap.tidb.tipb.SelectResponse;
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
+import io.grpc.stub.MetadataUtils;
 import io.prometheus.client.Histogram;
 import java.util.*;
 import java.util.function.Supplier;
@@ -43,7 +45,6 @@ import org.tikv.common.util.*;
 import org.tikv.kvproto.Coprocessor;
 import org.tikv.kvproto.Errorpb;
 import org.tikv.kvproto.Kvrpcpb.*;
-import org.tikv.kvproto.Metapb.Store;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
 import org.tikv.kvproto.TikvGrpc.TikvStub;
@@ -87,7 +88,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
   private RegionStoreClient(
       TiConfiguration conf,
       TiRegion region,
-      String storeVersion,
+      TiStore store,
       TiStoreType storeType,
       ChannelFactory channelFactory,
       TikvBlockingStub blockingStub,
@@ -95,15 +96,15 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       RegionManager regionManager,
       PDClient pdClient,
       RegionStoreClient.RegionStoreClientBuilder clientBuilder) {
-    super(conf, region, channelFactory, blockingStub, asyncStub, regionManager);
+    super(conf, region, store, channelFactory, blockingStub, asyncStub, regionManager);
     this.storeType = storeType;
 
     if (this.storeType == TiStoreType.TiKV) {
       this.lockResolverClient =
           AbstractLockResolverClient.getInstance(
-              storeVersion,
               conf,
               region,
+              store,
               this.blockingStub,
               this.asyncStub,
               channelFactory,
@@ -112,10 +113,10 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
               clientBuilder);
 
     } else {
-      Store tikvStore =
+      TiStore tikvStore =
           regionManager.getRegionStorePairByKey(region.getStartKey(), TiStoreType.TiKV).second;
 
-      String addressStr = tikvStore.getAddress();
+      String addressStr = tikvStore.getStore().getAddress();
       if (logger.isDebugEnabled()) {
         logger.debug(String.format("Create region store client on address %s", addressStr));
       }
@@ -126,9 +127,9 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
 
       this.lockResolverClient =
           AbstractLockResolverClient.getInstance(
-              tikvStore.getVersion(),
               conf,
               region,
+              tikvStore,
               tikvBlockingStub,
               tikvAsyncStub,
               channelFactory,
@@ -169,7 +170,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<GetRequest> factory =
         () ->
             GetRequest.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(version)))
+                .setContext(region.getReplicaContext(getResolvedLocks(version), this.storeType))
                 .setKey(key)
                 .setVersion(version)
                 .build();
@@ -214,7 +215,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<BatchGetRequest> request =
         () ->
             BatchGetRequest.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(version)))
+                .setContext(region.getReplicaContext(getResolvedLocks(version), this.storeType))
                 .addAllKeys(keys)
                 .setVersion(version)
                 .build();
@@ -277,7 +278,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<ScanRequest> request =
           () ->
               ScanRequest.newBuilder()
-                  .setContext(region.getContext(getResolvedLocks(version)))
+                  .setContext(region.getReplicaContext(getResolvedLocks(version), this.storeType))
                   .setStartKey(startKey)
                   .setVersion(version)
                   .setKeyOnly(keyOnly)
@@ -379,7 +380,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
           () ->
               getIsV4()
                   ? PrewriteRequest.newBuilder()
-                      .setContext(region.getContext())
+                      .setContext(region.getReplicaContext(storeType))
                       .setStartVersion(startTs)
                       .setPrimaryLock(primaryLock)
                       .addAllMutations(mutations)
@@ -389,7 +390,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
                       .setTxnSize(16)
                       .build()
                   : PrewriteRequest.newBuilder()
-                      .setContext(region.getContext())
+                      .setContext(region.getReplicaContext(storeType))
                       .setStartVersion(startTs)
                       .setPrimaryLock(primaryLock)
                       .addAllMutations(mutations)
@@ -469,7 +470,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<TxnHeartBeatRequest> factory =
           () ->
               TxnHeartBeatRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .setStartVersion(startTs)
                   .setPrimaryLock(primaryLock)
                   .setAdviseLockTtl(ttl)
@@ -527,7 +528,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
                 .setStartVersion(startTs)
                 .setCommitVersion(commitTs)
                 .addAllKeys(keys)
-                .setContext(region.getContext())
+                .setContext(region.getReplicaContext(storeType))
                 .build();
     KVErrorHandler<CommitResponse> handler =
         new KVErrorHandler<>(
@@ -588,7 +589,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<Coprocessor.Request> reqToSend =
         () ->
             Coprocessor.Request.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(startTs)))
+                .setContext(region.getReplicaContext(getResolvedLocks(startTs), this.storeType))
                 .setTp(REQ_TYPE_DAG.getValue())
                 .setStartTs(startTs)
                 .setData(req.toByteString())
@@ -711,7 +712,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<Coprocessor.Request> reqToSend =
         () ->
             Coprocessor.Request.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(startTs)))
+                .setContext(region.getReplicaContext(getResolvedLocks(startTs), this.storeType))
                 // TODO: If no executors...?
                 .setTp(REQ_TYPE_DAG.getValue())
                 .setData(req.toByteString())
@@ -749,7 +750,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<SplitRegionRequest> request =
         () ->
             SplitRegionRequest.newBuilder()
-                .setContext(region.getContext())
+                .setContext(region.getReplicaContext(storeType))
                 .addAllSplitKeys(splitKeys)
                 .build();
 
@@ -787,10 +788,11 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
                 new TiRegion(
                     region,
                     null,
+                    null,
                     conf.getIsolationLevel(),
                     conf.getCommandPriority(),
                     conf.getKvMode(),
-                    conf.isReplicaRead()))
+                    conf.getReplicaSelector()))
         .collect(Collectors.toList());
   }
 
@@ -801,7 +803,11 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
         GRPC_RAW_REQUEST_LATENCY.labels("client_grpc_raw_get").startTimer();
     try {
       Supplier<RawGetRequest> factory =
-          () -> RawGetRequest.newBuilder().setContext(region.getContext()).setKey(key).build();
+          () ->
+              RawGetRequest.newBuilder()
+                  .setContext(region.getReplicaContext(storeType))
+                  .setKey(key)
+                  .build();
       KVErrorHandler<RawGetResponse> handler =
           new KVErrorHandler<>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
@@ -833,7 +839,10 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     try {
       Supplier<RawGetKeyTTLRequest> factory =
           () ->
-              RawGetKeyTTLRequest.newBuilder().setContext(region.getContext()).setKey(key).build();
+              RawGetKeyTTLRequest.newBuilder()
+                  .setContext(region.getReplicaContext(storeType))
+                  .setKey(key)
+                  .build();
       KVErrorHandler<RawGetKeyTTLResponse> handler =
           new KVErrorHandler<>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
@@ -868,7 +877,11 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
         GRPC_RAW_REQUEST_LATENCY.labels("client_grpc_raw_delete").startTimer();
     try {
       Supplier<RawDeleteRequest> factory =
-          () -> RawDeleteRequest.newBuilder().setContext(region.getContext()).setKey(key).build();
+          () ->
+              RawDeleteRequest.newBuilder()
+                  .setContext(region.getReplicaContext(storeType))
+                  .setKey(key)
+                  .build();
 
       KVErrorHandler<RawDeleteResponse> handler =
           new KVErrorHandler<>(
@@ -902,7 +915,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawPutRequest> factory =
           () ->
               RawPutRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .setKey(key)
                   .setValue(value)
                   .setTtl(ttl)
@@ -940,7 +953,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawCASRequest> factory =
           () ->
               RawCASRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .setKey(key)
                   .setValue(value)
                   .setPreviousNotExist(true)
@@ -986,7 +999,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchGetRequest> factory =
           () ->
               RawBatchGetRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .addAllKeys(keys)
                   .build();
       KVErrorHandler<RawBatchGetResponse> handler =
@@ -1021,7 +1034,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchPutRequest> factory =
           () ->
               RawBatchPutRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .addAllPairs(kvPairs)
                   .setTtl(ttl)
                   .setForCas(atomic)
@@ -1073,7 +1086,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchDeleteRequest> factory =
           () ->
               RawBatchDeleteRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .addAllKeys(keys)
                   .setForCas(atomic)
                   .build();
@@ -1118,7 +1131,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawScanRequest> factory =
           () ->
               RawScanRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .setStartKey(key)
                   .setKeyOnly(keyOnly)
                   .setLimit(limit)
@@ -1164,7 +1177,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawDeleteRangeRequest> factory =
           () ->
               RawDeleteRangeRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getReplicaContext(storeType))
                   .setStartKey(startKey)
                   .setEndKey(endKey)
                   .build();
@@ -1232,25 +1245,48 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       this.pdClient = pdClient;
     }
 
-    public RegionStoreClient build(TiRegion region, Store store, TiStoreType storeType)
+    public RegionStoreClient build(TiRegion region, TiStore store, TiStoreType storeType)
         throws GrpcException {
       Objects.requireNonNull(region, "region is null");
       Objects.requireNonNull(store, "store is null");
       Objects.requireNonNull(storeType, "storeType is null");
 
-      String addressStr = store.getAddress();
+      String addressStr = store.getStore().getAddress();
       if (logger.isDebugEnabled()) {
         logger.debug(String.format("Create region store client on address %s", addressStr));
       }
-      ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+      ManagedChannel channel = null;
 
-      TikvBlockingStub blockingStub = TikvGrpc.newBlockingStub(channel);
+      TikvBlockingStub blockingStub = null;
-      TikvStub asyncStub = TikvGrpc.newStub(channel);
+      TikvStub asyncStub = null;
+
+      if (conf.getEnableGrpcForward() && region.getProxyStore() != null && store.isUnreachable()) {
+        addressStr = region.getProxyStore().getStore().getAddress();
+        channel =
+            channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
+        Metadata header = new Metadata();
+        header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
+        blockingStub = MetadataUtils.attachHeaders(TikvGrpc.newBlockingStub(channel), header);
+        asyncStub = MetadataUtils.attachHeaders(TikvGrpc.newStub(channel), header);
+      } else {
+        // If the store is reachable, which is update by check-health thread
+        if (!store.isUnreachable()) {
+          if (region.getProxyStore() != null) {
+            TiRegion newRegion = region.switchProxyStore(null);
+            if (regionManager.updateRegion(region, newRegion)) {
+              region = newRegion;
+            }
+          }
+        }
+        channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+        blockingStub = TikvGrpc.newBlockingStub(channel);
+        asyncStub = TikvGrpc.newStub(channel);
+      }
 
       return new RegionStoreClient(
           conf,
           region,
-          store.getVersion(),
+          store,
           storeType,
           channelFactory,
           blockingStub,
@@ -1260,7 +1296,8 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
           this);
     }
 
-    public synchronized RegionStoreClient build(TiRegion region, Store store) throws GrpcException {
+    public synchronized RegionStoreClient build(TiRegion region, TiStore store)
+        throws GrpcException {
       return build(region, store, TiStoreType.TiKV);
     }
 
@@ -1270,12 +1307,12 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
 
     public synchronized RegionStoreClient build(ByteString key, TiStoreType storeType)
         throws GrpcException {
-      Pair<TiRegion, Store> pair = regionManager.getRegionStorePairByKey(key, storeType);
+      Pair<TiRegion, TiStore> pair = regionManager.getRegionStorePairByKey(key, storeType);
       return build(pair.first, pair.second, storeType);
     }
 
     public synchronized RegionStoreClient build(TiRegion region) throws GrpcException {
-      Store store = regionManager.getStoreById(region.getLeader().getStoreId());
+      TiStore store = regionManager.getStoreById(region.getLeader().getStoreId());
       return build(region, store, TiStoreType.TiKV);
     }
 

src/main/java/org/tikv/common/region/TiRegion.java

@@ -22,14 +22,16 @@ import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
-import java.util.Random;
 import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.tikv.common.TiConfiguration.KVMode;
 import org.tikv.common.codec.Codec.BytesCodec;
 import org.tikv.common.codec.CodecDataInput;
 import org.tikv.common.codec.KeyUtils;
 import org.tikv.common.exception.TiClientInternalException;
 import org.tikv.common.key.Key;
+import org.tikv.common.replica.ReplicaSelector;
 import org.tikv.common.util.FastByteComparisons;
 import org.tikv.common.util.KeyRangeUtils;
 import org.tikv.kvproto.Kvrpcpb;
@@ -39,44 +41,33 @@ import org.tikv.kvproto.Metapb.Peer;
 import org.tikv.kvproto.Metapb.Region;
 
 public class TiRegion implements Serializable {
+  private static final Logger logger = LoggerFactory.getLogger(TiRegion.class);
+
   private final Region meta;
+  private final KVMode kvMode;
   private final IsolationLevel isolationLevel;
   private final Kvrpcpb.CommandPri commandPri;
   private final Peer leader;
-  private int followerIdx = 0;
+  private final ReplicaSelector replicaSelector;
-  private final boolean isReplicaRead;
+  private final List<Peer> replicaList;
-
+  private final TiStore proxyStore;
-  public TiRegion(
+  private int replicaIdx;
-      Region meta,
-      Peer leader,
-      IsolationLevel isolationLevel,
-      Kvrpcpb.CommandPri commandPri,
-      KVMode kvMode) {
-    this(meta, leader, isolationLevel, commandPri, kvMode, false);
-  }
-
-  private TiRegion(
-      Region meta,
-      Peer leader,
-      IsolationLevel isolationLevel,
-      Kvrpcpb.CommandPri commandPri,
-      boolean isReplicaRead) {
-    this.meta = meta;
-    this.leader = leader;
-    this.isolationLevel = isolationLevel;
-    this.commandPri = commandPri;
-    this.isReplicaRead = isReplicaRead;
-  }
 
   public TiRegion(
       Region meta,
       Peer leader,
+      TiStore proxyStore,
       IsolationLevel isolationLevel,
       Kvrpcpb.CommandPri commandPri,
       KVMode kvMode,
-      boolean isReplicaRead) {
+      ReplicaSelector replicaSelector) {
     Objects.requireNonNull(meta, "meta is null");
     this.meta = decodeRegion(meta, kvMode == KVMode.RAW);
+    this.kvMode = kvMode;
+    this.isolationLevel = isolationLevel;
+    this.commandPri = commandPri;
+    this.replicaSelector = replicaSelector;
+    this.proxyStore = proxyStore;
     if (leader == null || leader.getId() == 0) {
       if (meta.getPeersCount() == 0) {
         throw new TiClientInternalException("Empty peer list for region " + meta.getId());
@@ -86,17 +77,10 @@ public class TiRegion implements Serializable {
     } else {
       this.leader = leader;
     }
-    if (isReplicaRead && meta.getPeersCount() > 0) {
+
-      // try to get first follower
+    // init replicaList
-      try {
+    replicaList = replicaSelector.select(this.leader, getFollowerList(), getLearnerList());
-        chooseRandomFollower();
+    replicaIdx = 0;
-      } catch (Exception ignore) {
-        // ignore
-      }
-    }
-    this.isolationLevel = isolationLevel;
-    this.commandPri = commandPri;
-    this.isReplicaRead = isReplicaRead;
   }
 
   private Region decodeRegion(Region region, boolean isRawRegion) {
@@ -127,36 +111,41 @@ public class TiRegion implements Serializable {
     return leader;
   }
 
-  public Peer getCurrentFollower() {
+  public List<Peer> getFollowerList() {
-    return meta.getPeers(followerIdx);
+    List<Peer> peers = new ArrayList<>();
-  }
+    for (Peer peer : getMeta().getPeersList()) {
-
+      if (!peer.equals(this.leader)) {
-  private boolean isValidFollower(Peer peer) {
+        if (peer.getRole().equals(Metapb.PeerRole.Voter)) {
-    return Metapb.PeerRole.valueOf(peer.getRole().getValueDescriptor()) == Metapb.PeerRole.Voter;
+          peers.add(peer);
-  }
+        }
-
-  private void chooseRandomFollower() {
-    int cnt = meta.getPeersCount();
-    followerIdx = new Random().nextInt(cnt);
-    for (int retry = cnt - 1; retry > 0; retry--) {
-      followerIdx = (followerIdx + 1) % cnt;
-      Peer cur = meta.getPeers(followerIdx);
-      if (isValidFollower(cur)) {
-        return;
       }
     }
+    return peers;
   }
 
   public List<Peer> getLearnerList() {
     List<Peer> peers = new ArrayList<>();
     for (Peer peer : getMeta().getPeersList()) {
-      if (isValidFollower(peer)) {
+      if (peer.getRole().equals(Metapb.PeerRole.Learner)) {
         peers.add(peer);
       }
     }
     return peers;
   }
 
+  public Peer getCurrentReplica() {
+    return replicaList.get(replicaIdx);
+  }
+
+  public Peer getNextReplica() {
+    replicaIdx = (replicaIdx + 1) % replicaList.size();
+    return getCurrentReplica();
+  }
+
+  private boolean isLeader(Peer peer) {
+    return getLeader().equals(peer);
+  }
+
   public long getId() {
     return this.meta.getId();
   }
@@ -177,26 +166,30 @@ public class TiRegion implements Serializable {
     return Key.toRawKey(getEndKey());
   }
 
-  public Kvrpcpb.Context getContext() {
+  public Kvrpcpb.Context getLeaderContext() {
-    return getContext(java.util.Collections.emptySet());
+    return getContext(this.leader, java.util.Collections.emptySet(), TiStoreType.TiKV);
   }
 
-  public Kvrpcpb.Context getContext(Set<Long> resolvedLocks) {
+  public Kvrpcpb.Context getReplicaContext(TiStoreType storeType) {
+    return getContext(getCurrentReplica(), java.util.Collections.emptySet(), storeType);
+  }
+
+  public Kvrpcpb.Context getReplicaContext(Set<Long> resolvedLocks, TiStoreType storeType) {
+    return getContext(getCurrentReplica(), resolvedLocks, storeType);
+  }
+
+  private Kvrpcpb.Context getContext(
+      Peer currentPeer, Set<Long> resolvedLocks, TiStoreType storeType) {
+    boolean replicaRead = !isLeader(getCurrentReplica()) && TiStoreType.TiKV.equals(storeType);
+
     Kvrpcpb.Context.Builder builder = Kvrpcpb.Context.newBuilder();
-    builder.setIsolationLevel(this.isolationLevel);
+    builder
-    builder.setPriority(this.commandPri);
+        .setIsolationLevel(this.isolationLevel)
-    if (isReplicaRead) {
+        .setPriority(this.commandPri)
-      builder
+        .setRegionId(meta.getId())
-          .setRegionId(meta.getId())
+        .setPeer(currentPeer)
-          .setPeer(getCurrentFollower())
+        .setReplicaRead(replicaRead)
-          .setReplicaRead(true)
+        .setRegionEpoch(this.meta.getRegionEpoch());
-          .setRegionEpoch(this.meta.getRegionEpoch());
-    } else {
-      builder
-          .setRegionId(meta.getId())
-          .setPeer(this.leader)
-          .setRegionEpoch(this.meta.getRegionEpoch());
-    }
     builder.addAllResolvedLocks(resolvedLocks);
     return builder.build();
   }
@@ -207,6 +200,10 @@ public class TiRegion implements Serializable {
         meta.getId(), meta.getRegionEpoch().getConfVer(), meta.getRegionEpoch().getVersion());
   }
 
+  public TiStore getProxyStore() {
+    return proxyStore;
+  }
+
   /**
    * switches current peer to the one on specific store. It return false if no peer matches the
    * storeID.
@@ -218,12 +215,30 @@ public class TiRegion implements Serializable {
     List<Peer> peers = meta.getPeersList();
     for (Peer p : peers) {
       if (p.getStoreId() == leaderStoreID) {
-        return new TiRegion(this.meta, p, this.isolationLevel, this.commandPri, this.isReplicaRead);
+        return new TiRegion(
+            this.meta,
+            p,
+            this.proxyStore,
+            this.isolationLevel,
+            this.commandPri,
+            this.kvMode,
+            this.replicaSelector);
       }
     }
     return null;
   }
 
+  public TiRegion switchProxyStore(TiStore store) {
+    return new TiRegion(
+        this.meta,
+        this.leader,
+        store,
+        this.isolationLevel,
+        this.commandPri,
+        this.kvMode,
+        this.replicaSelector);
+  }
+
   public boolean isMoreThan(ByteString key) {
     return FastByteComparisons.compareTo(
             meta.getStartKey().toByteArray(),

src/main/java/org/tikv/common/region/TiStore.java

@@ -0,0 +1,34 @@
+package org.tikv.common.region;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.tikv.kvproto.Metapb;
+
+public class TiStore {
+  private final Metapb.Store store;
+  private AtomicBoolean unreachable;
+
+  public TiStore(Metapb.Store store) {
+    this.store = store;
+    this.unreachable = new AtomicBoolean(false);
+  }
+
+  public boolean markUnreachable() {
+    return this.unreachable.compareAndSet(false, true);
+  }
+
+  public void markReachable() {
+    this.unreachable.set(false);
+  }
+
+  public boolean isUnreachable() {
+    return this.unreachable.get();
+  }
+
+  public Metapb.Store getStore() {
+    return this.store;
+  }
+
+  public long getId() {
+    return this.store.getId();
+  }
+}

src/main/java/org/tikv/common/region/UnreachableStoreChecker.java

@@ -0,0 +1,84 @@
+package org.tikv.common.region;
+
+import io.grpc.ManagedChannel;
+import io.grpc.health.v1.HealthCheckRequest;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.LinkedBlockingQueue;
+import org.tikv.common.ReadOnlyPDClient;
+import org.tikv.common.util.ChannelFactory;
+import org.tikv.common.util.ConcreteBackOffer;
+import org.tikv.kvproto.Metapb;
+
+public class UnreachableStoreChecker implements Runnable {
+  private ConcurrentHashMap<Long, TiStore> stores;
+  private BlockingQueue<TiStore> taskQueue;
+  private final ChannelFactory channelFactory;
+  private final ReadOnlyPDClient pdClient;
+
+  public UnreachableStoreChecker(ChannelFactory channelFactory, ReadOnlyPDClient pdClient) {
+    this.stores = new ConcurrentHashMap();
+    this.taskQueue = new LinkedBlockingQueue<>();
+    this.channelFactory = channelFactory;
+    this.pdClient = pdClient;
+  }
+
+  public void scheduleStoreHealthCheck(TiStore store) {
+    TiStore oldStore = this.stores.get(Long.valueOf(store.getId()));
+    if (oldStore == store) {
+      return;
+    }
+    this.stores.put(Long.valueOf(store.getId()), store);
+    if (!this.taskQueue.add(store)) {
+      // add queue false, mark it reachable so that it can be put again.
+      store.markReachable();
+    }
+  }
+
+  private List<TiStore> getUnhealthStore() {
+    List<TiStore> unhealthStore = new LinkedList<>();
+    while (!this.taskQueue.isEmpty()) {
+      try {
+        TiStore store = this.taskQueue.take();
+        unhealthStore.add(store);
+      } catch (Exception e) {
+        return unhealthStore;
+      }
+    }
+    return unhealthStore;
+  }
+
+  @Override
+  public void run() {
+    List<TiStore> unhealthStore = getUnhealthStore();
+    for (TiStore store : unhealthStore) {
+      if (!store.isUnreachable()) {
+        continue;
+      }
+      String addressStr = store.getStore().getAddress();
+      ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+      HealthGrpc.HealthBlockingStub stub = HealthGrpc.newBlockingStub(channel);
+      HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
+      try {
+        HealthCheckResponse resp = stub.check(req);
+        if (resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING) {
+          store.markReachable();
+          this.stores.remove(Long.valueOf(store.getId()));
+          continue;
+        }
+        Metapb.Store newStore =
+            pdClient.getStore(ConcreteBackOffer.newRawKVBackOff(), store.getId());
+        if (newStore.getState() == Metapb.StoreState.Tombstone) {
+          continue;
+        }
+        this.taskQueue.add(store);
+      } catch (Exception e) {
+        this.taskQueue.add(store);
+      }
+    }
+  }
+}

src/main/java/org/tikv/common/replica/FollowerReplicaSelector.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tikv.common.replica;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.tikv.kvproto.Metapb;
+
+public class FollowerReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Metapb.Peer> select(
+      Metapb.Peer leader, List<Metapb.Peer> followers, List<Metapb.Peer> learners) {
+    List<Metapb.Peer> list = new ArrayList<>(followers);
+    Collections.shuffle(list);
+    return list;
+  }
+}

src/main/java/org/tikv/common/replica/LeaderFollowerReplicaSelector.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tikv.common.replica;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.tikv.kvproto.Metapb;
+
+public class LeaderFollowerReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Metapb.Peer> select(
+      Metapb.Peer leader, List<Metapb.Peer> followers, List<Metapb.Peer> learners) {
+    List<Metapb.Peer> list = new ArrayList<>(followers);
+    Collections.shuffle(list);
+    list.add(leader);
+    return list;
+  }
+}

src/main/java/org/tikv/common/replica/LeaderReplicaSelector.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tikv.common.replica;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.tikv.kvproto.Metapb;
+
+public class LeaderReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Metapb.Peer> select(
+      Metapb.Peer leader, List<Metapb.Peer> followers, List<Metapb.Peer> learners) {
+    List<Metapb.Peer> list = new ArrayList<>();
+    list.add(leader);
+    return list;
+  }
+}

src/main/java/org/tikv/common/replica/ReplicaSelector.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2021 PingCAP, Inc.
+ *
+ * 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,
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.tikv.common.replica;
+
+import java.io.Serializable;
+import java.util.List;
+import org.tikv.kvproto.Metapb;
+
+public interface ReplicaSelector extends Serializable {
+  public static final ReplicaSelector LEADER = new LeaderReplicaSelector();
+  public static final ReplicaSelector FOLLOWER = new FollowerReplicaSelector();
+  public static final ReplicaSelector LEADER_AND_FOLLOWER = new LeaderFollowerReplicaSelector();
+
+  List<Metapb.Peer> select(
+      Metapb.Peer leader, List<Metapb.Peer> followers, List<Metapb.Peer> learners);
+}

src/main/java/org/tikv/common/util/Batch.java

@@ -56,7 +56,7 @@ public class Batch {
   }
 
   public BackOffer getBackOffer() {
-    return backOffer;
+    return ConcreteBackOffer.create(backOffer);
   }
 
   public TiRegion getRegion() {

src/main/java/org/tikv/common/util/ClientUtils.java

@@ -121,7 +121,7 @@ public class ClientUtils {
 
   public static Map<TiRegion, List<ByteString>> groupKeysByRegion(
       RegionManager regionManager, List<ByteString> keys, BackOffer backoffer) {
-    return groupKeysByRegion(regionManager, new ArrayList<>(keys), backoffer, false);
+    return groupKeysByRegion(regionManager, keys, backoffer, false);
   }
 
   /**

src/main/java/org/tikv/common/util/RangeSplitter.java

@@ -29,9 +29,9 @@ import org.tikv.common.key.RowKey;
 import org.tikv.common.pd.PDUtils;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.region.TiStoreType;
 import org.tikv.kvproto.Coprocessor.KeyRange;
-import org.tikv.kvproto.Metapb;
 
 public class RangeSplitter {
   private final RegionManager regionManager;
@@ -51,12 +51,11 @@ public class RangeSplitter {
    * @param handles Handle list
    * @return <Region, HandleList> map
    */
-  public Map<Pair<TiRegion, Metapb.Store>, TLongArrayList> groupByAndSortHandlesByRegionId(
+  public Map<Pair<TiRegion, TiStore>, TLongArrayList> groupByAndSortHandlesByRegionId(
       long tableId, TLongArrayList handles) {
     TLongObjectHashMap<TLongArrayList> regionHandles = new TLongObjectHashMap<>();
-    TLongObjectHashMap<Pair<TiRegion, Metapb.Store>> idToRegionStorePair =
+    TLongObjectHashMap<Pair<TiRegion, TiStore>> idToRegionStorePair = new TLongObjectHashMap<>();
-        new TLongObjectHashMap<>();
+    Map<Pair<TiRegion, TiStore>, TLongArrayList> result = new HashMap<>();
-    Map<Pair<TiRegion, Metapb.Store>, TLongArrayList> result = new HashMap<>();
     handles.sort();
 
     byte[] endKey = null;
@@ -71,7 +70,7 @@ public class RangeSplitter {
           regionHandles.put(curRegion.getId(), handlesInCurRegion);
           handlesInCurRegion = new TLongArrayList();
         }
-        Pair<TiRegion, Metapb.Store> regionStorePair =
+        Pair<TiRegion, TiStore> regionStorePair =
             regionManager.getRegionStorePairByKey(ByteString.copyFrom(key.getBytes()));
         curRegion = regionStorePair.first;
         idToRegionStorePair.put(curRegion.getId(), regionStorePair);
@@ -84,7 +83,7 @@ public class RangeSplitter {
     }
     regionHandles.forEachEntry(
         (k, v) -> {
-          Pair<TiRegion, Metapb.Store> regionStorePair = idToRegionStorePair.get(k);
+          Pair<TiRegion, TiStore> regionStorePair = idToRegionStorePair.get(k);
           result.put(regionStorePair, v);
           return true;
         });
@@ -110,7 +109,7 @@ public class RangeSplitter {
     // Max value for current index handle range
     ImmutableList.Builder<RegionTask> regionTasks = ImmutableList.builder();
 
-    Map<Pair<TiRegion, Metapb.Store>, TLongArrayList> regionHandlesMap =
+    Map<Pair<TiRegion, TiStore>, TLongArrayList> regionHandlesMap =
         groupByAndSortHandlesByRegionId(tableId, handles);
 
     regionHandlesMap.forEach((k, v) -> createTask(0, v.size(), tableId, v, k, regionTasks));
@@ -123,7 +122,7 @@ public class RangeSplitter {
       int endPos,
       long tableId,
       TLongArrayList handles,
-      Pair<TiRegion, Metapb.Store> regionStorePair,
+      Pair<TiRegion, TiStore> regionStorePair,
       ImmutableList.Builder<RegionTask> regionTasks) {
     List<KeyRange> newKeyRanges = new ArrayList<>(endPos - startPos + 1);
     long startHandle = handles.get(startPos);
@@ -163,10 +162,10 @@ public class RangeSplitter {
     int i = 0;
     KeyRange range = keyRanges.get(i++);
     Map<Long, List<KeyRange>> idToRange = new HashMap<>(); // region id to keyRange list
-    Map<Long, Pair<TiRegion, Metapb.Store>> idToRegion = new HashMap<>();
+    Map<Long, Pair<TiRegion, TiStore>> idToRegion = new HashMap<>();
 
     while (true) {
-      Pair<TiRegion, Metapb.Store> regionStorePair =
+      Pair<TiRegion, TiStore> regionStorePair =
           regionManager.getRegionStorePairByKey(range.getStart(), storeType);
 
       if (regionStorePair == null) {
@@ -203,7 +202,7 @@ public class RangeSplitter {
     ImmutableList.Builder<RegionTask> resultBuilder = ImmutableList.builder();
     idToRange.forEach(
         (k, v) -> {
-          Pair<TiRegion, Metapb.Store> regionStorePair = idToRegion.get(k);
+          Pair<TiRegion, TiStore> regionStorePair = idToRegion.get(k);
           resultBuilder.add(new RegionTask(regionStorePair.first, regionStorePair.second, v));
         });
     return resultBuilder.build();
@@ -221,24 +220,23 @@ public class RangeSplitter {
 
   public static class RegionTask implements Serializable {
     private final TiRegion region;
-    private final Metapb.Store store;
+    private final TiStore store;
     private final List<KeyRange> ranges;
     private final String host;
 
-    RegionTask(TiRegion region, Metapb.Store store, List<KeyRange> ranges) {
+    RegionTask(TiRegion region, TiStore store, List<KeyRange> ranges) {
       this.region = region;
       this.store = store;
       this.ranges = ranges;
       String host = null;
       try {
-        host = PDUtils.addrToUri(store.getAddress()).getHost();
+        host = PDUtils.addrToUri(store.getStore().getAddress()).getHost();
       } catch (Exception ignored) {
       }
       this.host = host;
     }
 
-    public static RegionTask newInstance(
+    public static RegionTask newInstance(TiRegion region, TiStore store, List<KeyRange> ranges) {
-        TiRegion region, Metapb.Store store, List<KeyRange> ranges) {
       return new RegionTask(region, store, ranges);
     }
 
@@ -246,7 +244,7 @@ public class RangeSplitter {
       return region;
     }
 
-    public Metapb.Store getStore() {
+    public TiStore getStore() {
       return store;
     }
 

src/main/java/org/tikv/raw/RawKVClient.java

@@ -627,7 +627,7 @@ public class RawKVClient implements AutoCloseable {
     } catch (final TiKVException e) {
       // TODO: any elegant way to re-split the ranges if fails?
       backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
-      logger.warn("ReSplitting ranges for BatchPutRequest");
+      logger.debug("ReSplitting ranges for BatchPutRequest");
       // retry
       return doSendBatchPutWithRefetchRegion(backOffer, batch);
     }
@@ -685,7 +685,7 @@ public class RawKVClient implements AutoCloseable {
     } catch (final TiKVException e) {
       backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
       clientBuilder.getRegionManager().invalidateRegion(batch.getRegion());
-      logger.warn("ReSplitting ranges for BatchGetRequest", e);
+      logger.debug("ReSplitting ranges for BatchGetRequest", e);
 
       // retry
       return Pair.create(doSendBatchGetWithRefetchRegion(backOffer, batch), new ArrayList<>());
@@ -726,7 +726,7 @@ public class RawKVClient implements AutoCloseable {
     } catch (final TiKVException e) {
       backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
       clientBuilder.getRegionManager().invalidateRegion(batch.getRegion());
-      logger.warn("ReSplitting ranges for BatchGetRequest", e);
+      logger.debug("ReSplitting ranges for BatchGetRequest", e);
 
       // retry
       return doSendBatchDeleteWithRefetchRegion(backOffer, batch);
@@ -776,7 +776,7 @@ public class RawKVClient implements AutoCloseable {
     } catch (final TiKVException e) {
       backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
       clientBuilder.getRegionManager().invalidateRegion(range.getRegion());
-      logger.warn("ReSplitting ranges for BatchDeleteRangeRequest", e);
+      logger.debug("ReSplitting ranges for BatchDeleteRangeRequest", e);
 
       // retry
       return doSendDeleteRangeWithRefetchRegion(backOffer, range);
@@ -810,7 +810,9 @@ public class RawKVClient implements AutoCloseable {
   private List<TiRegion> fetchRegionsFromRange(
       BackOffer backOffer, ByteString startKey, ByteString endKey) {
     List<TiRegion> regions = new ArrayList<>();
-    while (startKey.isEmpty() || Key.toRawKey(startKey).compareTo(Key.toRawKey(endKey)) < 0) {
+    while (startKey.isEmpty()
+        || endKey.isEmpty()
+        || Key.toRawKey(startKey).compareTo(Key.toRawKey(endKey)) < 0) {
       TiRegion currentRegion = clientBuilder.getRegionManager().getRegionByKey(startKey, backOffer);
       regions.add(currentRegion);
       startKey = currentRegion.getEndKey();

src/main/java/org/tikv/txn/AbstractLockResolverClient.java

@@ -26,6 +26,7 @@ import org.tikv.common.exception.KeyException;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
 import org.tikv.kvproto.Kvrpcpb;
@@ -66,22 +67,23 @@ public interface AbstractLockResolverClient {
   }
 
   static AbstractLockResolverClient getInstance(
-      String storeVersion,
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvGrpc.TikvBlockingStub blockingStub,
       TikvGrpc.TikvStub asyncStub,
       ChannelFactory channelFactory,
       RegionManager regionManager,
       PDClient pdClient,
       RegionStoreClient.RegionStoreClientBuilder clientBuilder) {
-    if (StoreVersion.compareTo(storeVersion, Version.RESOLVE_LOCK_V3) < 0) {
+    if (StoreVersion.compareTo(store.getStore().getVersion(), Version.RESOLVE_LOCK_V3) < 0) {
       return new LockResolverClientV2(
-          conf, region, blockingStub, asyncStub, channelFactory, regionManager);
+          conf, region, store, blockingStub, asyncStub, channelFactory, regionManager);
-    } else if (StoreVersion.compareTo(storeVersion, Version.RESOLVE_LOCK_V4) < 0) {
+    } else if (StoreVersion.compareTo(store.getStore().getVersion(), Version.RESOLVE_LOCK_V4) < 0) {
       return new LockResolverClientV3(
           conf,
           region,
+          store,
           blockingStub,
           asyncStub,
           channelFactory,
@@ -92,6 +94,7 @@ public interface AbstractLockResolverClient {
       return new LockResolverClientV4(
           conf,
           region,
+          store,
           blockingStub,
           asyncStub,
           channelFactory,

src/main/java/org/tikv/txn/BatchKeys.java

@@ -1,62 +0,0 @@
-/*
- * Copyright 2017 PingCAP, Inc.
- *
- * 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,
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.tikv.txn;
-
-import com.google.protobuf.ByteString;
-import java.util.ArrayList;
-import java.util.List;
-import org.tikv.common.region.TiRegion;
-import org.tikv.kvproto.Metapb;
-
-public class BatchKeys {
-  private final TiRegion region;
-  private final Metapb.Store store;
-  private List<ByteString> keys;
-  private final int sizeInBytes;
-
-  public BatchKeys(
-      TiRegion region, Metapb.Store store, List<ByteString> keysInput, int sizeInBytes) {
-    this.region = region;
-    this.store = store;
-    this.keys = new ArrayList<>();
-    this.keys.addAll(keysInput);
-    this.sizeInBytes = sizeInBytes;
-  }
-
-  public List<ByteString> getKeys() {
-    return keys;
-  }
-
-  public void setKeys(List<ByteString> keys) {
-    this.keys = keys;
-  }
-
-  public TiRegion getRegion() {
-    return region;
-  }
-
-  public Metapb.Store getStore() {
-    return store;
-  }
-
-  public int getSizeInBytes() {
-    return sizeInBytes;
-  }
-
-  public float getSizeInKB() {
-    return ((float) sizeInBytes) / 1024;
-  }
-}

src/main/java/org/tikv/txn/ClientRPCResult.java

@@ -1,52 +0,0 @@
-/*
- * Copyright 2019 The 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,
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.tikv.txn;
-
-public class ClientRPCResult {
-  private boolean success;
-  private boolean retry;
-  private Exception exception;
-
-  public ClientRPCResult(boolean success, boolean retry, Exception exception) {
-    this.success = success;
-    this.retry = retry;
-    this.exception = exception;
-  }
-
-  public boolean isSuccess() {
-    return success;
-  }
-
-  public void setSuccess(boolean success) {
-    this.success = success;
-  }
-
-  public boolean isRetry() {
-    return retry;
-  }
-
-  public void setRetry(boolean retry) {
-    this.retry = retry;
-  }
-
-  public Exception getException() {
-    return exception;
-  }
-
-  public void setException(Exception exception) {
-    this.exception = exception;
-  }
-}

src/main/java/org/tikv/txn/GroupKeyResult.java

@@ -1,41 +0,0 @@
-/*
- * Copyright 2017 PingCAP, Inc.
- *
- * 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,
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.tikv.txn;
-
-import com.google.protobuf.ByteString;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import org.tikv.common.region.TiRegion;
-import org.tikv.common.util.Pair;
-import org.tikv.kvproto.Metapb;
-
-public class GroupKeyResult {
-
-  private Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupsResult;
-
-  public GroupKeyResult() {
-    this.groupsResult = new HashMap<>();
-  }
-
-  public Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> getGroupsResult() {
-    return groupsResult;
-  }
-
-  public void setGroupsResult(Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupsResult) {
-    this.groupsResult = groupsResult;
-  }
-}

src/main/java/org/tikv/txn/LockResolverClientV2.java

@@ -42,6 +42,7 @@ import org.tikv.common.region.AbstractRegionStoreClient;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.TiRegion;
 import org.tikv.common.region.TiRegion.RegionVerID;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
 import org.tikv.common.util.TsoUtils;
@@ -74,11 +75,12 @@ public class LockResolverClientV2 extends AbstractRegionStoreClient
   public LockResolverClientV2(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
       TikvStub asyncStub,
       ChannelFactory channelFactory,
       RegionManager regionManager) {
-    super(conf, region, channelFactory, blockingStub, asyncStub, regionManager);
+    super(conf, region, store, channelFactory, blockingStub, asyncStub, regionManager);
     resolved = new HashMap<>();
     recentResolved = new LinkedList<>();
     readWriteLock = new ReentrantReadWriteLock();
@@ -125,7 +127,7 @@ public class LockResolverClientV2 extends AbstractRegionStoreClient
       Supplier<CleanupRequest> factory =
           () ->
               CleanupRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getLeaderContext())
                   .setKey(primary)
                   .setStartVersion(txnID)
                   .build();
@@ -232,7 +234,7 @@ public class LockResolverClientV2 extends AbstractRegionStoreClient
         factory =
             () ->
                 ResolveLockRequest.newBuilder()
-                    .setContext(region.getContext())
+                    .setContext(region.getLeaderContext())
                     .setStartVersion(lock.getTxnID())
                     .setCommitVersion(txnStatus)
                     .build();
@@ -240,7 +242,7 @@ public class LockResolverClientV2 extends AbstractRegionStoreClient
         factory =
             () ->
                 ResolveLockRequest.newBuilder()
-                    .setContext(region.getContext())
+                    .setContext(region.getLeaderContext())
                     .setStartVersion(lock.getTxnID())
                     .build();
       }

src/main/java/org/tikv/txn/LockResolverClientV3.java

@@ -39,10 +39,7 @@ import org.tikv.common.exception.KeyException;
 import org.tikv.common.exception.RegionException;
 import org.tikv.common.exception.TiClientInternalException;
 import org.tikv.common.operation.KVErrorHandler;
-import org.tikv.common.region.AbstractRegionStoreClient;
+import org.tikv.common.region.*;
-import org.tikv.common.region.RegionManager;
-import org.tikv.common.region.RegionStoreClient;
-import org.tikv.common.region.TiRegion;
 import org.tikv.common.region.TiRegion.RegionVerID;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
@@ -79,13 +76,14 @@ public class LockResolverClientV3 extends AbstractRegionStoreClient
   public LockResolverClientV3(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
       TikvStub asyncStub,
       ChannelFactory channelFactory,
       RegionManager regionManager,
       PDClient pdClient,
       RegionStoreClient.RegionStoreClientBuilder clientBuilder) {
-    super(conf, region, channelFactory, blockingStub, asyncStub, regionManager);
+    super(conf, region, store, channelFactory, blockingStub, asyncStub, regionManager);
     resolved = new HashMap<>();
     recentResolved = new LinkedList<>();
     readWriteLock = new ReentrantReadWriteLock();
@@ -151,7 +149,7 @@ public class LockResolverClientV3 extends AbstractRegionStoreClient
 
       Kvrpcpb.ResolveLockRequest.Builder builder =
           Kvrpcpb.ResolveLockRequest.newBuilder()
-              .setContext(region.getContext())
+              .setContext(region.getLeaderContext())
               .setStartVersion(lock.getTxnID());
 
       if (txnStatus.isCommitted()) {
@@ -230,7 +228,7 @@ public class LockResolverClientV3 extends AbstractRegionStoreClient
         () -> {
           TiRegion primaryKeyRegion = regionManager.getRegionByKey(primary);
           return CleanupRequest.newBuilder()
-              .setContext(primaryKeyRegion.getContext())
+              .setContext(primaryKeyRegion.getLeaderContext())
               .setKey(primary)
               .setStartVersion(txnID)
               .setCurrentTs(currentTS)

src/main/java/org/tikv/txn/LockResolverClientV4.java

@@ -39,10 +39,7 @@ import org.tikv.common.exception.KeyException;
 import org.tikv.common.exception.RegionException;
 import org.tikv.common.exception.TiClientInternalException;
 import org.tikv.common.operation.KVErrorHandler;
-import org.tikv.common.region.AbstractRegionStoreClient;
+import org.tikv.common.region.*;
-import org.tikv.common.region.RegionManager;
-import org.tikv.common.region.RegionStoreClient;
-import org.tikv.common.region.TiRegion;
 import org.tikv.common.region.TiRegion.RegionVerID;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
@@ -79,13 +76,14 @@ public class LockResolverClientV4 extends AbstractRegionStoreClient
   public LockResolverClientV4(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
       TikvStub asyncStub,
       ChannelFactory channelFactory,
       RegionManager regionManager,
       PDClient pdClient,
       RegionStoreClient.RegionStoreClientBuilder clientBuilder) {
-    super(conf, region, channelFactory, blockingStub, asyncStub, regionManager);
+    super(conf, region, store, channelFactory, blockingStub, asyncStub, regionManager);
     resolved = new HashMap<>();
     recentResolved = new LinkedList<>();
     readWriteLock = new ReentrantReadWriteLock();
@@ -169,7 +167,7 @@ public class LockResolverClientV4 extends AbstractRegionStoreClient
       Supplier<Kvrpcpb.PessimisticRollbackRequest> factory =
           () ->
               Kvrpcpb.PessimisticRollbackRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(region.getLeaderContext())
                   .setStartVersion(lock.getTxnID())
                   .setForUpdateTs(forUpdateTS)
                   .addKeys(lock.getKey())
@@ -287,7 +285,7 @@ public class LockResolverClientV4 extends AbstractRegionStoreClient
         () -> {
           TiRegion primaryKeyRegion = regionManager.getRegionByKey(primary);
           return Kvrpcpb.CheckTxnStatusRequest.newBuilder()
-              .setContext(primaryKeyRegion.getContext())
+              .setContext(primaryKeyRegion.getLeaderContext())
               .setPrimaryKey(primary)
               .setLockTs(txnID)
               .setCallerStartTs(callerStartTS)
@@ -364,7 +362,7 @@ public class LockResolverClientV4 extends AbstractRegionStoreClient
 
       Kvrpcpb.ResolveLockRequest.Builder builder =
           Kvrpcpb.ResolveLockRequest.newBuilder()
-              .setContext(region.getContext())
+              .setContext(region.getLeaderContext())
               .setStartVersion(lock.getTxnID());
 
       if (txnStatus.isCommitted()) {

src/main/java/org/tikv/txn/TTLManager.java

@@ -30,11 +30,11 @@ import org.tikv.common.exception.TiBatchWriteException;
 import org.tikv.common.meta.TiTimestamp;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffFunction;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.Pair;
-import org.tikv.kvproto.Metapb;
 import org.tikv.txn.type.ClientRPCResult;
 
 /**
@@ -105,9 +105,9 @@ public class TTLManager {
   }
 
   private void sendTxnHeartBeat(BackOffer bo, long ttl) {
-    Pair<TiRegion, Metapb.Store> pair = regionManager.getRegionStorePairByKey(primaryLock);
+    Pair<TiRegion, TiStore> pair = regionManager.getRegionStorePairByKey(primaryLock);
     TiRegion tiRegion = pair.first;
-    Metapb.Store store = pair.second;
+    TiStore store = pair.second;
 
     ClientRPCResult result = kvClient.txnHeartBeat(bo, primaryLock, startTS, ttl, tiRegion, store);
 
@@ -121,7 +121,7 @@ public class TTLManager {
             new GrpcException(
                 String.format("sendTxnHeartBeat failed, regionId=%s", tiRegion.getId()),
                 result.getException()));
-        this.regionManager.invalidateStore(store.getId());
+        this.regionManager.invalidateStore(store.getStore().getId());
         this.regionManager.invalidateRegion(tiRegion);
         // re-split keys and commit again.
         sendTxnHeartBeat(bo, ttl);

src/main/java/org/tikv/txn/TwoPhaseCommitter.java

@@ -38,13 +38,13 @@ import org.tikv.common.exception.GrpcException;
 import org.tikv.common.exception.TiBatchWriteException;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffFunction;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.Pair;
 import org.tikv.kvproto.Kvrpcpb;
 import org.tikv.kvproto.Kvrpcpb.Op;
-import org.tikv.kvproto.Metapb;
 import org.tikv.txn.type.BatchKeys;
 import org.tikv.txn.type.ClientRPCResult;
 import org.tikv.txn.type.GroupKeyResult;
@@ -146,9 +146,9 @@ public class TwoPhaseCommitter {
 
   private void doPrewritePrimaryKeyWithRetry(BackOffer backOffer, ByteString key, ByteString value)
       throws TiBatchWriteException {
-    Pair<TiRegion, Metapb.Store> pair = this.regionManager.getRegionStorePairByKey(key, backOffer);
+    Pair<TiRegion, TiStore> pair = this.regionManager.getRegionStorePairByKey(key, backOffer);
     TiRegion tiRegion = pair.first;
-    Metapb.Store store = pair.second;
+    TiStore store = pair.second;
 
     Kvrpcpb.Mutation mutation;
     if (!value.isEmpty()) {
@@ -201,9 +201,9 @@ public class TwoPhaseCommitter {
 
   private void doCommitPrimaryKeyWithRetry(BackOffer backOffer, ByteString key, long commitTs)
       throws TiBatchWriteException {
-    Pair<TiRegion, Metapb.Store> pair = this.regionManager.getRegionStorePairByKey(key, backOffer);
+    Pair<TiRegion, TiStore> pair = this.regionManager.getRegionStorePairByKey(key, backOffer);
     TiRegion tiRegion = pair.first;
-    Metapb.Store store = pair.second;
+    TiStore store = pair.second;
     ByteString[] keys = new ByteString[] {key};
 
     // send rpc request to tikv server
@@ -335,11 +335,11 @@ public class TwoPhaseCommitter {
     // groups keys by region
     GroupKeyResult groupResult = this.groupKeysByRegion(keys, size, backOffer);
     List<BatchKeys> batchKeyList = new LinkedList<>();
-    Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupKeyMap = groupResult.getGroupsResult();
+    Map<Pair<TiRegion, TiStore>, List<ByteString>> groupKeyMap = groupResult.getGroupsResult();
 
-    for (Map.Entry<Pair<TiRegion, Metapb.Store>, List<ByteString>> entry : groupKeyMap.entrySet()) {
+    for (Map.Entry<Pair<TiRegion, TiStore>, List<ByteString>> entry : groupKeyMap.entrySet()) {
       TiRegion tiRegion = entry.getKey().first;
-      Metapb.Store store = entry.getKey().second;
+      TiStore store = entry.getKey().second;
       this.appendBatchBySize(batchKeyList, tiRegion, store, entry.getValue(), true, mutations);
     }
 
@@ -450,7 +450,7 @@ public class TwoPhaseCommitter {
   private void appendBatchBySize(
       List<BatchKeys> batchKeyList,
       TiRegion tiRegion,
-      Metapb.Store store,
+      TiStore store,
       List<ByteString> keys,
       boolean sizeIncludeValue,
       Map<ByteString, Kvrpcpb.Mutation> mutations) {
@@ -571,11 +571,11 @@ public class TwoPhaseCommitter {
     // groups keys by region
     GroupKeyResult groupResult = this.groupKeysByRegion(keys, size, backOffer);
     List<BatchKeys> batchKeyList = new ArrayList<>();
-    Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupKeyMap = groupResult.getGroupsResult();
+    Map<Pair<TiRegion, TiStore>, List<ByteString>> groupKeyMap = groupResult.getGroupsResult();
 
-    for (Map.Entry<Pair<TiRegion, Metapb.Store>, List<ByteString>> entry : groupKeyMap.entrySet()) {
+    for (Map.Entry<Pair<TiRegion, TiStore>, List<ByteString>> entry : groupKeyMap.entrySet()) {
       TiRegion tiRegion = entry.getKey().first;
-      Metapb.Store store = entry.getKey().second;
+      TiStore store = entry.getKey().second;
       this.appendBatchBySize(batchKeyList, tiRegion, store, entry.getValue(), false, null);
     }
 
@@ -615,13 +615,12 @@ public class TwoPhaseCommitter {
 
   private GroupKeyResult groupKeysByRegion(ByteString[] keys, int size, BackOffer backOffer)
       throws TiBatchWriteException {
-    Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groups = new HashMap<>();
+    Map<Pair<TiRegion, TiStore>, List<ByteString>> groups = new HashMap<>();
     int index = 0;
     try {
       for (; index < size; index++) {
         ByteString key = keys[index];
-        Pair<TiRegion, Metapb.Store> pair =
+        Pair<TiRegion, TiStore> pair = this.regionManager.getRegionStorePairByKey(key, backOffer);
-            this.regionManager.getRegionStorePairByKey(key, backOffer);
         if (pair != null) {
           groups.computeIfAbsent(pair, e -> new ArrayList<>()).add(key);
         }

src/main/java/org/tikv/txn/TxnKVClient.java

@@ -33,11 +33,11 @@ import org.tikv.common.meta.TiTimestamp;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffFunction;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.kvproto.Kvrpcpb;
-import org.tikv.kvproto.Metapb;
 import org.tikv.txn.type.ClientRPCResult;
 
 /** KV client of transaction APIs for GET/PUT/DELETE/SCAN */
@@ -94,7 +94,7 @@ public class TxnKVClient implements AutoCloseable {
       long lockTTL,
       long startTs,
       TiRegion tiRegion,
-      Metapb.Store store) {
+      TiStore store) {
     ClientRPCResult result = new ClientRPCResult(true, false, null);
     // send request
     RegionStoreClient client = clientBuilder.build(tiRegion, store);
@@ -116,7 +116,7 @@ public class TxnKVClient implements AutoCloseable {
       long startTs,
       long ttl,
       TiRegion tiRegion,
-      Metapb.Store store) {
+      TiStore store) {
     ClientRPCResult result = new ClientRPCResult(true, false, null);
     // send request
     RegionStoreClient client = clientBuilder.build(tiRegion, store);
@@ -148,7 +148,7 @@ public class TxnKVClient implements AutoCloseable {
       long startTs,
       long commitTs,
       TiRegion tiRegion,
-      Metapb.Store store) {
+      TiStore store) {
     ClientRPCResult result = new ClientRPCResult(true, false, null);
     // send request
     RegionStoreClient client = clientBuilder.build(tiRegion, store);

src/main/java/org/tikv/txn/type/BatchKeys.java

@@ -19,16 +19,15 @@ import com.google.protobuf.ByteString;
 import java.util.ArrayList;
 import java.util.List;
 import org.tikv.common.region.TiRegion;
-import org.tikv.kvproto.Metapb;
+import org.tikv.common.region.TiStore;
 
 public class BatchKeys {
   private final TiRegion region;
-  private final Metapb.Store store;
+  private final TiStore store;
   private List<ByteString> keys;
   private final int sizeInBytes;
 
-  public BatchKeys(
+  public BatchKeys(TiRegion region, TiStore store, List<ByteString> keysInput, int sizeInBytes) {
-      TiRegion region, Metapb.Store store, List<ByteString> keysInput, int sizeInBytes) {
     this.region = region;
     this.store = store;
     this.keys = new ArrayList<>();
@@ -48,7 +47,7 @@ public class BatchKeys {
     return region;
   }
 
-  public Metapb.Store getStore() {
+  public TiStore getStore() {
     return store;
   }
 

src/main/java/org/tikv/txn/type/GroupKeyResult.java

@@ -20,22 +20,22 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.Pair;
-import org.tikv.kvproto.Metapb;
 
 public class GroupKeyResult {
 
-  private Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupsResult;
+  private Map<Pair<TiRegion, TiStore>, List<ByteString>> groupsResult;
 
   public GroupKeyResult() {
     this.groupsResult = new HashMap<>();
   }
 
-  public Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> getGroupsResult() {
+  public Map<Pair<TiRegion, TiStore>, List<ByteString>> getGroupsResult() {
     return groupsResult;
   }
 
-  public void setGroupsResult(Map<Pair<TiRegion, Metapb.Store>, List<ByteString>> groupsResult) {
+  public void setGroupsResult(Map<Pair<TiRegion, TiStore>, List<ByteString>> groupsResult) {
     this.groupsResult = groupsResult;
   }
 }

src/test/java/org/tikv/BUILD

@@ -1,24 +0,0 @@
-package(default_visibility = ["//visibility:public"])
-
-load(":rule.bzl", "junit_suite_test")
-
-junit_suite_test(
-    name = "tikv-client-java-test",
-    srcs = glob(
-        ["**/*.java"],
-    ),
-    deps = [
-        "//src/main/java/com/pingcap/tikv:tikv-java-client-lib",
-        "//:java",
-        "//:java_compile_imports",
-        "@com_fasterxml_jackson_core_jackson_annotations//jar",
-        "@com_fasterxml_jackson_core_jackson_core//jar",
-        "@com_fasterxml_jackson_core_jackson_databind//jar",
-
-        "@org_pubref_rules_protobuf//java:grpc_compiletime_deps",
-        "@org_pubref_rules_protobuf//java:netty_runtime_deps",
-        "@net_sf_trove4j_trove4j//jar",
-        "@junit_junit//jar",
-        "@joda_time//jar",
-    ],
-)

src/test/java/org/tikv/common/MockServerTest.java

@@ -5,6 +5,7 @@ import java.io.IOException;
 import org.junit.Before;
 import org.tikv.common.TiConfiguration.KVMode;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.replica.ReplicaSelector;
 import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.Pdpb;
 
@@ -31,9 +32,11 @@ public class MockServerTest extends PDMockServerTest {
         new TiRegion(
             r,
             r.getPeers(0),
+            null,
             session.getConf().getIsolationLevel(),
             session.getConf().getCommandPriority(),
-            KVMode.TXN);
+            KVMode.TXN,
+            ReplicaSelector.LEADER);
     pdServer.addGetRegionResp(Pdpb.GetRegionResponse.newBuilder().setRegion(r).build());
     server = new KVMockServer();
     port = server.start(region);

src/test/java/org/tikv/common/PDClientTest.java

@@ -18,7 +18,6 @@ package org.tikv.common;
 import static org.junit.Assert.*;
 import static org.tikv.common.GrpcUtils.encodeKey;
 
-import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.util.concurrent.*;
 import org.junit.Test;
@@ -37,7 +36,7 @@ public class PDClientTest extends PDMockServerTest {
   @Test
   public void testCreate() throws Exception {
     try (PDClient client = session.getPDClient()) {
-      assertEquals(client.getLeaderWrapper().getLeaderInfo(), LOCAL_ADDR + ":" + pdServer.port);
+      assertEquals(client.getPdClientWrapper().getLeaderInfo(), LOCAL_ADDR + ":" + pdServer.port);
       assertEquals(client.getHeader().getClusterId(), CLUSTER_ID);
     }
   }
@@ -45,17 +44,18 @@ public class PDClientTest extends PDMockServerTest {
   @Test
   public void testSwitchLeader() throws Exception {
     try (PDClient client = session.getPDClient()) {
-      client.switchLeader(ImmutableList.of("http://" + LOCAL_ADDR + ":" + (pdServer.port + 1)));
+      client.trySwitchLeader("http://" + LOCAL_ADDR + ":" + (pdServer.port + 1));
       assertEquals(
-          client.getLeaderWrapper().getLeaderInfo(), LOCAL_ADDR + ":" + (pdServer.port + 1));
+          "http://" + LOCAL_ADDR + ":" + (pdServer.port + 1),
+          client.getPdClientWrapper().getLeaderInfo());
     }
     tearDown();
     setUp(LOCAL_ADDR_IPV6);
     try (PDClient client = session.getPDClient()) {
-      client.switchLeader(
+      client.trySwitchLeader("http://" + LOCAL_ADDR_IPV6 + ":" + (pdServer.port + 2));
-          ImmutableList.of("http://" + LOCAL_ADDR_IPV6 + ":" + (pdServer.port + 2)));
       assertEquals(
-          client.getLeaderWrapper().getLeaderInfo(), LOCAL_ADDR_IPV6 + ":" + (pdServer.port + 2));
+          "http://" + LOCAL_ADDR_IPV6 + ":" + (pdServer.port + 2),
+          client.getPdClientWrapper().getLeaderInfo());
     }
   }
 
@@ -95,33 +95,6 @@ public class PDClientTest extends PDMockServerTest {
     }
   }
 
-  @Test
-  public void testGetRegionByKeyAsync() throws Exception {
-    byte[] startKey = new byte[] {1, 0, 2, 4};
-    byte[] endKey = new byte[] {1, 0, 2, 5};
-    int confVer = 1026;
-    int ver = 1027;
-    pdServer.addGetRegionResp(
-        GrpcUtils.makeGetRegionResponse(
-            pdServer.getClusterId(),
-            GrpcUtils.makeRegion(
-                1,
-                encodeKey(startKey),
-                encodeKey(endKey),
-                GrpcUtils.makeRegionEpoch(confVer, ver),
-                GrpcUtils.makePeer(1, 10),
-                GrpcUtils.makePeer(2, 20))));
-    try (PDClient client = session.getPDClient()) {
-      TiRegion r = client.getRegionByKeyAsync(defaultBackOff(), ByteString.EMPTY).get();
-      assertEquals(r.getStartKey(), ByteString.copyFrom(startKey));
-      assertEquals(r.getEndKey(), ByteString.copyFrom(endKey));
-      assertEquals(r.getRegionEpoch().getConfVer(), confVer);
-      assertEquals(r.getRegionEpoch().getVersion(), ver);
-      assertEquals(r.getLeader().getId(), 1);
-      assertEquals(r.getLeader().getStoreId(), 10);
-    }
-  }
-
   @Test
   public void testGetRegionById() throws Exception {
     byte[] startKey = new byte[] {1, 0, 2, 4};
@@ -150,33 +123,6 @@ public class PDClientTest extends PDMockServerTest {
     }
   }
 
-  @Test
-  public void testGetRegionByIdAsync() throws Exception {
-    byte[] startKey = new byte[] {1, 0, 2, 4};
-    byte[] endKey = new byte[] {1, 0, 2, 5};
-    int confVer = 1026;
-    int ver = 1027;
-    pdServer.addGetRegionByIDResp(
-        GrpcUtils.makeGetRegionResponse(
-            pdServer.getClusterId(),
-            GrpcUtils.makeRegion(
-                1,
-                encodeKey(startKey),
-                encodeKey(endKey),
-                GrpcUtils.makeRegionEpoch(confVer, ver),
-                GrpcUtils.makePeer(1, 10),
-                GrpcUtils.makePeer(2, 20))));
-    try (PDClient client = session.getPDClient()) {
-      TiRegion r = client.getRegionByIDAsync(defaultBackOff(), 0).get();
-      assertEquals(r.getStartKey(), ByteString.copyFrom(startKey));
-      assertEquals(r.getEndKey(), ByteString.copyFrom(endKey));
-      assertEquals(r.getRegionEpoch().getConfVer(), confVer);
-      assertEquals(r.getRegionEpoch().getVersion(), ver);
-      assertEquals(r.getLeader().getId(), 1);
-      assertEquals(r.getLeader().getStoreId(), 10);
-    }
-  }
-
   @Test
   public void testGetStore() throws Exception {
     long storeId = 1;
@@ -208,38 +154,6 @@ public class PDClientTest extends PDMockServerTest {
     }
   }
 
-  @Test
-  public void testGetStoreAsync() throws Exception {
-    long storeId = 1;
-    String testAddress = "testAddress";
-    pdServer.addGetStoreResp(
-        GrpcUtils.makeGetStoreResponse(
-            pdServer.getClusterId(),
-            GrpcUtils.makeStore(
-                storeId,
-                testAddress,
-                Metapb.StoreState.Up,
-                GrpcUtils.makeStoreLabel("k1", "v1"),
-                GrpcUtils.makeStoreLabel("k2", "v2"))));
-    try (PDClient client = session.getPDClient()) {
-      Store r = client.getStoreAsync(defaultBackOff(), 0).get();
-      assertEquals(r.getId(), storeId);
-      assertEquals(r.getAddress(), testAddress);
-      assertEquals(r.getState(), Metapb.StoreState.Up);
-      assertEquals(r.getLabels(0).getKey(), "k1");
-      assertEquals(r.getLabels(1).getKey(), "k2");
-      assertEquals(r.getLabels(0).getValue(), "v1");
-      assertEquals(r.getLabels(1).getValue(), "v2");
-
-      pdServer.addGetStoreResp(
-          GrpcUtils.makeGetStoreResponse(
-              pdServer.getClusterId(),
-              GrpcUtils.makeStore(storeId, testAddress, Metapb.StoreState.Tombstone)));
-      assertEquals(
-          StoreState.Tombstone, client.getStoreAsync(defaultBackOff(), 0).get().getState());
-    }
-  }
-
   private BackOffer defaultBackOff() {
     return ConcreteBackOffer.newCustomBackOff(1000);
   }

src/test/java/org/tikv/common/PDMockServer.java

@@ -43,7 +43,7 @@ public class PDMockServer extends PDGrpc.PDImplBase {
   @Override
   public void getMembers(GetMembersRequest request, StreamObserver<GetMembersResponse> resp) {
     try {
-      resp.onNext(getMembersResp.removeFirst().get());
+      resp.onNext(getMembersResp.getFirst().get());
       resp.onCompleted();
     } catch (Exception e) {
       resp.onError(Status.INTERNAL.asRuntimeException());

src/test/java/org/tikv/common/RegionManagerTest.java

@@ -26,10 +26,10 @@ import org.junit.Test;
 import org.tikv.common.key.Key;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.KeyRangeUtils;
 import org.tikv.common.util.Pair;
 import org.tikv.kvproto.Metapb;
-import org.tikv.kvproto.Metapb.Store;
 import org.tikv.kvproto.Metapb.StoreState;
 
 public class RegionManagerTest extends PDMockServerTest {
@@ -115,7 +115,7 @@ public class RegionManagerTest extends PDMockServerTest {
                 Metapb.StoreState.Up,
                 GrpcUtils.makeStoreLabel("k1", "v1"),
                 GrpcUtils.makeStoreLabel("k2", "v2"))));
-    Pair<TiRegion, Store> pair = mgr.getRegionStorePairByKey(searchKey);
+    Pair<TiRegion, TiStore> pair = mgr.getRegionStorePairByKey(searchKey);
     assertEquals(pair.first.getId(), regionId);
     assertEquals(pair.first.getId(), storeId);
   }
@@ -133,8 +133,8 @@ public class RegionManagerTest extends PDMockServerTest {
                 Metapb.StoreState.Up,
                 GrpcUtils.makeStoreLabel("k1", "v1"),
                 GrpcUtils.makeStoreLabel("k2", "v2"))));
-    Store store = mgr.getStoreById(storeId);
+    TiStore store = mgr.getStoreById(storeId);
-    assertEquals(store.getId(), storeId);
+    assertEquals(store.getStore().getId(), storeId);
 
     pdServer.addGetStoreResp(
         GrpcUtils.makeGetStoreResponse(

src/test/java/org/tikv/common/RegionStoreClientTest.java

@@ -24,6 +24,7 @@ import org.junit.Test;
 import org.tikv.common.region.RegionManager;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
+import org.tikv.common.region.TiStore;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.kvproto.Kvrpcpb;
@@ -40,13 +41,14 @@ public class RegionStoreClientTest extends MockServerTest {
   }
 
   private RegionStoreClient createClient(String version) {
-    Metapb.Store store =
+    Metapb.Store meta =
         Metapb.Store.newBuilder()
             .setAddress(LOCAL_ADDR + ":" + port)
             .setId(1)
             .setState(Metapb.StoreState.Up)
             .setVersion(version)
             .build();
+    TiStore store = new TiStore(meta);
 
     RegionStoreClientBuilder builder =
         new RegionStoreClientBuilder(

src/test/java/org/tikv/raw/RawKVClientTest.java

@@ -1,5 +1,7 @@
 package org.tikv.raw;
 
+import static org.junit.Assert.*;
+
 import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.util.*;
@@ -8,6 +10,7 @@ import java.util.stream.Collectors;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -91,7 +94,8 @@ public class RawKVClientTest {
     }
   }
 
-  @Test
+  // tikv-4.0 does not support atomic api
+  @Ignore
   public void atomicAPITest() {
     if (!initialized) return;
     long ttl = 10;
@@ -100,19 +104,20 @@ public class RawKVClientTest {
     ByteString value2 = ByteString.copyFromUtf8("value2");
     client.delete(key);
     ByteString res1 = client.putIfAbsent(key, value, ttl);
-    assert res1.isEmpty();
+    assertTrue(res1.isEmpty());
     ByteString res2 = client.putIfAbsent(key, value2, ttl);
-    assert res2.equals(value);
+    assertEquals(value, res2);
     try {
       Thread.sleep(ttl * 1000);
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
     }
     ByteString res3 = client.putIfAbsent(key, value, ttl);
-    assert res3.isEmpty();
+    assertTrue(res3.isEmpty());
   }
 
-  @Test
+  // tikv-4.0 doest not support ttl
+  @Ignore
   public void getKeyTTLTest() {
     if (!initialized) return;
     long ttl = 10;
@@ -254,30 +259,47 @@ public class RawKVClientTest {
   public void simpleTest() {
     if (!initialized) return;
     ByteString key = rawKey("key");
+    ByteString key0 = rawKey("key0");
     ByteString key1 = rawKey("key1");
     ByteString key2 = rawKey("key2");
     ByteString key3 = rawKey("key3");
+    ByteString value = rawValue("value");
     ByteString value1 = rawValue("value1");
     ByteString value2 = rawValue("value2");
+    ByteString value3 = rawValue("value3");
+    Kvrpcpb.KvPair kv = Kvrpcpb.KvPair.newBuilder().setKey(key).setValue(value).build();
     Kvrpcpb.KvPair kv1 = Kvrpcpb.KvPair.newBuilder().setKey(key1).setValue(value1).build();
     Kvrpcpb.KvPair kv2 = Kvrpcpb.KvPair.newBuilder().setKey(key2).setValue(value2).build();
+    Kvrpcpb.KvPair kv3 = Kvrpcpb.KvPair.newBuilder().setKey(key3).setValue(value3).build();
 
     try {
-      checkEmpty(key1);
+      checkDeleteRange(ByteString.EMPTY, ByteString.EMPTY);
-      checkEmpty(key2);
+      checkEmpty(kv);
-      checkPut(key1, value1);
+      checkEmpty(kv1);
-      checkPut(key2, value2);
+      checkEmpty(kv2);
-      List<Kvrpcpb.KvPair> result = new ArrayList<>();
+      checkEmpty(kv3);
-      List<Kvrpcpb.KvPair> result2 = new ArrayList<>();
+      checkPut(kv);
-      result.add(kv1);
+      checkPut(kv1);
-      result.add(kv2);
+      checkPut(kv2);
-      checkScan(key, key3, result, limit);
+      checkPut(kv3);
-      checkScan(key1, key3, result, limit);
+      // <key, value>, <key1,value1>, <key2,value2>, <key3,value3>
-      checkScan(key, key1, new ArrayList<>(), limit);
+      // (-∞, +∞)
-      result2.add(kv1);
+      checkScan(ByteString.EMPTY, ByteString.EMPTY, Arrays.asList(kv, kv1, kv2, kv3), limit);
-      checkScan(key, key2, result2, limit);
+      // (-∞, key3)
+      checkScan(ByteString.EMPTY, key3, Arrays.asList(kv, kv1, kv2), limit);
+      // [key1, +∞)
+      checkScan(key1, ByteString.EMPTY, Arrays.asList(kv1, kv2, kv3), limit);
+      // [key, key3)
+      checkScan(key, key3, Arrays.asList(kv, kv1, kv2), limit);
+      // [key1, key3)
+      checkScan(key1, key3, Arrays.asList(kv1, kv2), limit);
+      // [key0, key1)
+      checkScan(key0, key1, new ArrayList<>(), limit);
+      // [key, key2)
+      checkScan(key, key2, Arrays.asList(kv, kv1), limit);
       checkDelete(key1);
       checkDelete(key2);
+      checkDeleteRange(ByteString.EMPTY, ByteString.EMPTY);
     } catch (final TiKVException e) {
       logger.warn("Test fails with Exception: " + e);
     }
@@ -509,7 +531,7 @@ public class RawKVClientTest {
     } else {
       int i = 0;
       for (Map.Entry<ByteString, ByteString> pair : data.entrySet()) {
-        assert client.get(pair.getKey()).equals(pair.getValue());
+        assertEquals(pair.getValue(), client.get(pair.getKey()));
         i++;
         if (i >= getCases) {
           break;
@@ -756,27 +778,31 @@ public class RawKVClientTest {
   private void checkBatchGet(List<ByteString> keys) {
     List<Kvrpcpb.KvPair> result = client.batchGet(keys);
     for (Kvrpcpb.KvPair kvPair : result) {
-      assert data.containsKey(kvPair.getKey());
+      assertTrue(data.containsKey(kvPair.getKey()));
-      assert kvPair.getValue().equals(data.get(kvPair.getKey()));
+      assertEquals(data.get(kvPair.getKey()), kvPair.getValue());
     }
   }
 
+  private void checkPut(Kvrpcpb.KvPair kv) {
+    checkPut(kv.getKey(), kv.getValue());
+  }
+
   private void checkPut(ByteString key, ByteString value) {
     client.put(key, value);
-    assert client.get(key).equals(value);
+    assertEquals(value, client.get(key));
   }
 
   private void checkBatchPut(Map<ByteString, ByteString> kvPairs) {
     client.batchPut(kvPairs);
     for (Map.Entry<ByteString, ByteString> kvPair : kvPairs.entrySet()) {
-      assert client.get(kvPair.getKey()).equals(kvPair.getValue());
+      assertEquals(kvPair.getValue(), client.get(kvPair.getKey()));
     }
   }
 
   private void checkScan(
-      ByteString startKey, ByteString endKey, List<Kvrpcpb.KvPair> ans, int limit) {
+      ByteString startKey, ByteString endKey, List<Kvrpcpb.KvPair> expected, int limit) {
     List<Kvrpcpb.KvPair> result = client.scan(startKey, endKey, limit);
-    assert result.equals(ans);
+    assertEquals(expected, result);
   }
 
   private void checkScan(
@@ -812,7 +838,7 @@ public class RawKVClientTest {
                           .setValue(kvPair.getValue())
                           .build())
               .collect(Collectors.toList());
-      assert result.get(i).equals(partialResult);
+      assertEquals(partialResult, result.get(i));
       i++;
     }
   }
@@ -827,31 +853,35 @@ public class RawKVClientTest {
     logger.info("delete range complete");
     List<Kvrpcpb.KvPair> result = client.scan(startKey, endKey);
     logger.info("checking scan complete. number of remaining keys in range: " + result.size());
-    assert result.isEmpty();
+    assertTrue(result.isEmpty());
   }
 
   private void checkPutTTL(ByteString key, ByteString value, long ttl) {
     client.put(key, value, ttl);
-    assert client.get(key).equals(value);
+    assertEquals(value, client.get(key));
   }
 
   private void checkGetKeyTTL(ByteString key, long ttl) {
     Long t = client.getKeyTTL(key);
-    assert t != null;
+    assertNotNull(t);
-    assert t <= ttl && t > 0;
+    assertTrue(t <= ttl && t > 0);
   }
 
   private void checkGetTTLTimeOut(ByteString key) {
-    assert client.get(key).isEmpty();
+    assertTrue(client.get(key).isEmpty());
   }
 
   private void checkGetKeyTTLTimeOut(ByteString key) {
     Long t = client.getKeyTTL(key);
-    assert t == null;
+    assertNull(t);
+  }
+
+  private void checkEmpty(Kvrpcpb.KvPair kv) {
+    checkEmpty(kv.getKey());
   }
 
   private void checkEmpty(ByteString key) {
-    assert client.get(key).isEmpty();
+    assertTrue(client.get(key).isEmpty());
   }
 
   private static ByteString rawKey(String key) {

src/test/java/org/tikv/rule.bzl

@@ -1,40 +0,0 @@
-def junit_suite_test(name, srcs, deps, size="small", resources=[], classpath_resources=[], jvm_flags=[], tags=[], data=[]):
-  tests = []
-  package = PACKAGE_NAME.replace("src/test/java/", "").replace("/", ".")
-  for src in srcs:
-    if src.endswith("Test.java"):
-      if "/" in src:
-         src = package + "." + src.replace("/", ".")
-      tests += [src.replace(".java", ".class")]
-
-
-  native.genrule(
-    name = name + "-AllTests-gen",
-    outs = ["AllTests.java"],
-    cmd = """
-      cat <<EOF >> $@
-package %s;
-
-import org.junit.runner.RunWith;
-import org.junit.runners.Suite;
-
-@RunWith(Suite.class)
-@Suite.SuiteClasses({%s})
-public class AllTests {}
-EOF
-    """ % (package, ",".join(tests))
-  )
-
-  native.java_test(
-    name = name,
-    srcs = srcs + ["AllTests.java"],
-    test_class = package + ".AllTests",
-    resources = resources,
-    classpath_resources = classpath_resources,
-    data = data,
-    size = size,
-    tags = tags,
-    jvm_flags = jvm_flags,
-    deps = deps + [
-    ],
-  )

src/test/java/org/tikv/txn/ReplicaReadTest.java

@@ -0,0 +1,84 @@
+package org.tikv.txn;
+
+import com.google.protobuf.ByteString;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.tikv.common.TiConfiguration;
+import org.tikv.common.TiSession;
+import org.tikv.common.replica.ReplicaSelector;
+import org.tikv.kvproto.Metapb;
+
+public class ReplicaReadTest extends TXNTest {
+  private TiSession session;
+  private String key;
+  private String value;
+
+  @Test
+  public void leaderReadTest() {
+    doTest(TiConfiguration.ReplicaRead.LEADER);
+  }
+
+  // ci only has one TiKV instance
+  @Ignore
+  public void followerReadTest() {
+    doTest(TiConfiguration.ReplicaRead.FOLLOWER);
+  }
+
+  @Test
+  public void leadAndFollowerReadTest() {
+    doTest(TiConfiguration.ReplicaRead.LEADER_AND_FOLLOWER);
+  }
+
+  @Test
+  public void replicaSelectorTest() {
+    TiConfiguration conf = TiConfiguration.createDefault();
+
+    conf.setReplicaSelector(
+        new ReplicaSelector() {
+          @Override
+          public List<Metapb.Peer> select(
+              Metapb.Peer leader, List<Metapb.Peer> followers, List<Metapb.Peer> learners) {
+            List<Metapb.Peer> list = new ArrayList<>();
+            list.addAll(followers);
+            list.addAll(learners);
+            list.add(leader);
+            return list;
+          }
+        });
+    session = TiSession.create(conf);
+
+    putKV(key, value);
+    ByteString v = session.createSnapshot().get(ByteString.copyFromUtf8(key));
+    Assert.assertEquals(value, v.toStringUtf8());
+  }
+
+  private void doTest(TiConfiguration.ReplicaRead replicaRead) {
+    TiConfiguration conf = TiConfiguration.createDefault();
+    conf.setReplicaRead(replicaRead);
+    session = TiSession.create(conf);
+
+    putKV(key, value);
+    ByteString v = session.createSnapshot().get(ByteString.copyFromUtf8(key));
+    Assert.assertEquals(value, v.toStringUtf8());
+  }
+
+  @Before
+  public void setUp() {
+    super.setUp();
+    key = genRandomKey(64);
+    value = "v0";
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    if (session != null) {
+      session.close();
+    }
+    super.tearDown();
+  }
+}

src/test/java/org/tikv/txn/TXNTest.java

@@ -0,0 +1,131 @@
+package org.tikv.txn;
+
+import static junit.framework.TestCase.assertTrue;
+import static junit.framework.TestCase.fail;
+
+import com.google.protobuf.ByteString;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import org.junit.After;
+import org.junit.Before;
+import org.tikv.common.TiConfiguration;
+import org.tikv.common.TiSession;
+import org.tikv.common.exception.RegionException;
+import org.tikv.common.region.RegionStoreClient;
+import org.tikv.common.region.TiRegion;
+import org.tikv.common.util.BackOffFunction;
+import org.tikv.common.util.BackOffer;
+import org.tikv.common.util.ConcreteBackOffer;
+import org.tikv.kvproto.Kvrpcpb;
+
+public class TXNTest {
+  static final int DEFAULT_TTL = 10;
+  private TiSession session;
+  RegionStoreClient.RegionStoreClientBuilder builder;
+
+  @Before
+  public void setUp() {
+    TiConfiguration conf = TiConfiguration.createDefault();
+    try {
+      session = TiSession.create(conf);
+      this.builder = session.getRegionStoreClientBuilder();
+    } catch (Exception e) {
+      fail("TiDB cluster may not be present");
+    }
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    if (session != null) {
+      session.close();
+    }
+  }
+
+  void putKV(String key, String value) {
+    long startTS = session.getTimestamp().getVersion();
+    long commitTS = session.getTimestamp().getVersion();
+    putKV(key, value, startTS, commitTS);
+  }
+
+  void putKV(String key, String value, long startTS, long commitTS) {
+    Kvrpcpb.Mutation m =
+        Kvrpcpb.Mutation.newBuilder()
+            .setKey(ByteString.copyFromUtf8(key))
+            .setOp(Kvrpcpb.Op.Put)
+            .setValue(ByteString.copyFromUtf8(value))
+            .build();
+
+    boolean res = prewriteString(Collections.singletonList(m), startTS, key, DEFAULT_TTL);
+    assertTrue(res);
+    res = commitString(Collections.singletonList(key), startTS, commitTS);
+    assertTrue(res);
+  }
+
+  boolean prewriteString(List<Kvrpcpb.Mutation> mutations, long startTS, String primary, long ttl) {
+    return prewrite(mutations, startTS, ByteString.copyFromUtf8(primary), ttl);
+  }
+
+  boolean prewrite(List<Kvrpcpb.Mutation> mutations, long startTS, ByteString primary, long ttl) {
+    if (mutations.size() == 0) return true;
+    BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(1000);
+
+    for (Kvrpcpb.Mutation m : mutations) {
+      while (true) {
+        try {
+          TiRegion region = session.getRegionManager().getRegionByKey(m.getKey());
+          RegionStoreClient client = builder.build(region);
+          client.prewrite(backOffer, primary, Collections.singletonList(m), startTS, ttl, false);
+          break;
+        } catch (RegionException e) {
+          backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+        }
+      }
+    }
+    return true;
+  }
+
+  boolean commitString(List<String> keys, long startTS, long commitTS) {
+    return commit(
+        keys.stream().map(ByteString::copyFromUtf8).collect(Collectors.toList()),
+        startTS,
+        commitTS);
+  }
+
+  boolean commit(List<ByteString> keys, long startTS, long commitTS) {
+    if (keys.size() == 0) return true;
+    BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(1000);
+
+    for (ByteString byteStringK : keys) {
+      while (true) {
+        try {
+          TiRegion tiRegion = session.getRegionManager().getRegionByKey(byteStringK);
+          RegionStoreClient client = builder.build(tiRegion);
+          client.commit(backOffer, Collections.singletonList(byteStringK), startTS, commitTS);
+          break;
+        } catch (RegionException e) {
+          backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+        }
+      }
+    }
+    return true;
+  }
+
+  String genRandomKey(int strLength) {
+    Random rnd = ThreadLocalRandom.current();
+    String prefix = rnd.nextInt(2) % 2 == 0 ? "a-test-" : "z-test-";
+    StringBuilder ret = new StringBuilder(prefix);
+    for (int i = 0; i < strLength; i++) {
+      boolean isChar = (rnd.nextInt(2) % 2 == 0);
+      if (isChar) {
+        int choice = rnd.nextInt(2) % 2 == 0 ? 65 : 97;
+        ret.append((char) (choice + rnd.nextInt(26)));
+      } else {
+        ret.append(rnd.nextInt(10));
+      }
+    }
+    return ret.toString();
+  }
+}