update google-gson to 2.8.9

Signed-off-by: marsishandsome <marsishandsome@gmail.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/labeler.yml

@@ -0,0 +1,17 @@
+component/config:
+  - config/*
+
+component/br:
+  - src/main/java/org/tikv/br/*
+
+component/cdc:
+  - src/main/java/org/tikv/cdc/*
+
+component/common:
+  - src/main/java/org/tikv/common/*
+
+component/rawkv:
+  - src/main/java/org/tikv/raw/*
+
+component/txnkv:
+  - src/main/java/org/tikv/txn/*

.github/pr-branch-labeler.yml

@@ -0,0 +1,5 @@
+type/3.1-cherry-pick:
+  base: "release-3.1"
+
+type/3.0-cherry-pick:
+  base: "release-3.0"

.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

.github/workflows/labeler.yml

@@ -0,0 +1,11 @@
+name: "Pull Request Labeler"
+on:
+- pull_request_target
+
+jobs:
+  triage:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/labeler@v3
+      with:
+        repo-token: "${{ secrets.GITHUB_TOKEN }}"

.github/workflows/pr-branch-labeler.yml

@@ -0,0 +1,13 @@
+name: PR Branch Labeler
+
+on: pull_request
+
+jobs:
+  label_prs:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Label PRs
+      if: github.event.action == 'opened' # Only run the action when the PR was first opened
+      uses: ffittschen/pr-branch-labeler@v1
+      with:
+        repo-token: ${{ secrets.GITHUB_TOKEN }}

Makefile

@@ -0,0 +1,15 @@
+.DEFAULT: build
+
+.PHONY: clean
+
+build:
+	mvn clean package -Dmaven.test.skip=true
+
+fmt:
+	./dev/javafmt
+
+test:
+	mvn clean test
+
+clean:
+	mvn clean

README.md

@@ -27,7 +27,7 @@ The jar can be found in `./target/`
 
 ## Usage
 
-This project is designed to hook with `pd` and `tikv` which you can find in `PingCAP` github page.
+This project is designed to hook with `[pd](https://github.com/tikv/pd)` and `[tikv](https://github.com/tikv/tikv)`.
 
 When you work with this project, you have to communicate with `pd` and `tikv`. Please run TiKV and PD in advance.
 
@@ -56,7 +56,7 @@ After building, add following lines into your `pom.xml` if you are using Maven
 <dependency>
 	<groupId>org.tikv</groupId>
 	<artifactId>tikv-client-java</artifactId>
-	<version>3.0.0</version>
+	<version>3.1.0</version>
 </dependency>
 ```
 
@@ -66,6 +66,7 @@ After building, add following lines into your `pom.xml` if you are using Maven
 ### Create a RawKVClient
 
 ```java
+import org.tikv.common.TiConfiguration;
 import org.tikv.common.TiSession;
 import org.tikv.raw.RawKVClient;
 
@@ -74,109 +75,124 @@ public class Main {
 		// You MUST create a raw configuration if you are using RawKVClient.
 		TiConfiguration conf = TiConfiguration.createRawDefault(YOUR_PD_ADDRESSES);
 		TiSession session = TiSession.create(conf);
-		RawKVClient = session.createRawKVClient();
+		RawKVClient client = session.createRawClient();
 	}
 }
 ```
 
-### API
+## Java Client Configuration Parameter
 
-```java
-/**
- * Put a raw key-value pair to TiKV
- *
- * @param key raw key
- * @param value raw value
- */
-void put(ByteString key, ByteString value)
+### JVM Parameter
+
+The following includes JVM related parameters.
+
+#### tikv.pd.addresses
+- pd addresses, separated by comma
+- default: 127.0.0.1:2379
+
+#### tikv.grpc.timeout_in_ms
+- timeout of grpc request  
+- default: 600ms
+
+#### tikv.grpc.scan_timeout_in_ms
+- timeout of scan/delete range grpc request
+- default: 20s
+
+### Metrics Parameter
+
+#### tikv.metrics.enable
+- whether to enable metrics exporting
+- default: false
+
+#### tikv.metrics.port
+- the metrics exporting http port
+- default: 3140
+
+### ThreadPool Parameter
+
+The following includes ThreadPool related parameters, which can be passed in through JVM parameters.
+
+#### tikv.batch_get_concurrency
+- the thread pool size of batchGet on client side
+- default: 20
+
+#### tikv.batch_put_concurrency
+- the thread pool size of batchPut on client side
+- default: 20
+
+#### tikv.batch_delete_concurrency
+- the thread pool size of batchDelete on client side
+- default: 20
+
+#### tikv.batch_scan_concurrency
+- the thread pool size of batchScan on client side
+- default: 5
+
+#### tikv.delete_range_concurrency
+- the thread pool size of deleteRange on client side
+- default: 20
+
+#### tikv.rawkv.default_backoff_in_ms
+- RawKV default backoff in milliseconds
+- default: 20000 (20 seconds)
+
+#### tikv.rawkv.read_timeout_in_ms
+- RawKV read timeout in milliseconds. This parameter controls the timeout of `get` `getKeyTTL`.
+- default: 2000 (2 seconds)
+
+#### tikv.rawkv.write_timeout_in_ms
+- RawKV write timeout in milliseconds. This parameter controls the timeout of `put` `putAtomic` `putIfAbsent` `delete` `deleteAtomic`.
+- default: 2000 (2 seconds)
+
+#### tikv.rawkv.batch_read_timeout_in_ms
+- RawKV batch read timeout in milliseconds. This parameter controls the timeout of `batchGet`.
+- default: 2000 (2 seconds)
+
+#### tikv.rawkv.batch_write_timeout_in_ms
+- RawKV batch write timeout in milliseconds. This parameter controls the timeout of `batchPut` `batchDelete` `batchDeleteAtomic`.
+- default: 2000 (2 seconds)
+
+#### tikv.rawkv.scan_timeout_in_ms
+- RawKV scan timeout in milliseconds. This parameter controls the timeout of `batchScan` `scan` `scanPrefix`.
+- default: 10000 (10 seconds)
+
+#### tikv.rawkv.clean_timeout_in_ms
+- RawKV clean timeout in milliseconds. This parameter controls the timeout of `deleteRange` `deletePrefix`.
+- default: 600000 (10 minutes)
+
+## Metrics
+
+Client Java supports exporting metrics to Prometheus using poll mode and viewing on Grafana. The following steps shows how to enable this function.
+
+### Step 1: Enable metrics exporting
+
+- set the config `tikv.metrics.enable` to `true`
+- call TiConfiguration.setMetricsEnable(true)
+
+### Step 2: Set the metrics port
+
+- set the config `tikv.metrics.port`
+- call TiConfiguration.setMetricsPort
+
+Default port is 3140.
+
+### Step 3: Config Prometheus
+
+Add the following config to `conf/prometheus.yml` and restart Prometheus.
+
+```yaml
+- job_name: "tikv-client"
+    honor_labels: true
+    static_configs:
+    - targets:
+        - '127.0.0.1:3140'
+        - '127.0.0.2:3140'
+        - '127.0.0.3:3140'
 ```
 
-```java
-/**
- * Get a raw key-value pair from TiKV if key exists
- *
- * @param key raw key
- * @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
- */
-ByteString get(ByteString key)
-```
-
-```java
-/**
- * Scan raw key-value pairs from TiKV in range [startKey, endKey)
- *
- * @param startKey raw start key, inclusive
- * @param endKey raw end key, exclusive
- * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
- * @return list of key-value pairs in range
- */
-List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, int limit)
-```
-
-```java
-/**
- * Scan raw key-value pairs from TiKV in range [startKey, endKey)
- *
- * @param startKey raw start key, inclusive
- * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
- * @return list of key-value pairs in range
- */
-List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit)
-```
-
-```java
-/**
- * Delete a raw key-value pair from TiKV if key exists
- *
- * @param key raw key to be deleted
- */
-void delete(ByteString key)
-```
-
-## Java Client 配置参数
-
-本文介绍了与部署使用 Java Client 相关的配置参数。
-
-### 常用配置 JVM 参数
-
-以下包括常用配置的 JVM 相关参数。
-
-####tikv.pd.addresses
-- pd 集群的地址，逗号分隔
-- 默认值 127.0.0.1:2379
-
-####tikv.grpc.timeout_in_ms
-- grpc 请求的 timeout 时间
-- 默认值 600ms
-
-####tikv.grpc.scan_timeout_in_ms
-- scan/delete range grpc 请求的 timeout 时间
-- 默认值 20s
-
-### ThreadPool 配置 JVM 参数
-
-以下包括 ThreadPool 相关的参数及其默认配置，可通过 JVM 参数传入。
-
-####tikv.batch_get_concurrency
-- Client 端 batchGet 请求的线程池大小
-- 默认值 20
-
-####tikv.batch_put_concurrency
-- Client 端 batchPut 请求的线程池大小
-- 默认值 20
-
-####tikv.batch_delete_concurrency
-- Client 端 batchDelete 请求的线程池大小
-- 默认值 20
-
-####tikv.batch_scan_concurrency
-- Client 端 batchScan 请求的线程池大小
-- 默认值 5
-
-####tikv.delete_range_concurrency
-- Client 端 deleteRange 请求的线程池大小
-- 默认值 20
+### Step 4: Config Grafana
 
+Import the [Client-Java-Summary dashboard config](/metrics/grafana/client_java_summary.json) to Grafana.
 
 ## License
-Apache 2.0 license. See the [LICENSE](./LICENSE) file for details.
+Apache 2.0 license. See the [LICENSE](./LICENSE) file for details.

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.11.1-jh</version>
     <packaging>jar</packaging>
     <name>TiKV Java Client</name>
     <description>A Java Client for TiKV</description>
@@ -62,9 +62,11 @@
         <protobuf.version>3.5.1</protobuf.version>
         <log4j.version>1.2.17</log4j.version>
         <slf4j.version>1.7.16</slf4j.version>
-        <grpc.version>1.24.0</grpc.version>
+        <grpc.version>1.38.0</grpc.version>
+        <gson.version>2.8.9</gson.version>
         <powermock.version>1.6.6</powermock.version>
-        <jackson.version>2.10.0</jackson.version>
+        <jackson.version>2.10.5</jackson.version>
+        <jackson.databind.version>2.10.5.1</jackson.databind.version>
         <trove4j.version>3.0.1</trove4j.version>
         <jetcd.version>0.4.1</jetcd.version>
         <joda-time.version>2.9.9</joda-time.version>
@@ -126,12 +128,22 @@
             <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>
             <version>${grpc.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.google.code.gson</groupId>
+            <artifactId>gson</artifactId>
+            <version>${gson.version}</version>
+        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>
@@ -140,7 +152,7 @@
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-databind</artifactId>
-            <version>${jackson.version}</version>
+            <version>${jackson.databind.version}</version>
         </dependency>
         <dependency>
             <groupId>io.etcd</groupId>
@@ -180,7 +192,7 @@
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-lang3</artifactId>
-            <version>3.9</version>
+            <version>3.10</version>
             <scope>test</scope>
         </dependency>
         <dependency>

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,16 @@ 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.AbstractFutureStub;
 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;
@@ -33,14 +39,15 @@ import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
 
 public abstract class AbstractGRPCClient<
-        BlockingStubT extends AbstractStub<BlockingStubT>, StubT extends AbstractStub<StubT>>
+        BlockingStubT extends AbstractStub<BlockingStubT>,
+        FutureStubT extends AbstractFutureStub<FutureStubT>>
     implements AutoCloseable {
   protected final Logger logger = LoggerFactory.getLogger(this.getClass());
   protected final ChannelFactory channelFactory;
   protected TiConfiguration conf;
   protected long timeout;
   protected BlockingStubT blockingStub;
-  protected StubT asyncStub;
+  protected FutureStubT asyncStub;
 
   protected AbstractGRPCClient(TiConfiguration conf, ChannelFactory channelFactory) {
     this.conf = conf;
@@ -52,7 +59,7 @@ public abstract class AbstractGRPCClient<
       TiConfiguration conf,
       ChannelFactory channelFactory,
       BlockingStubT blockingStub,
-      StubT asyncStub) {
+      FutureStubT asyncStub) {
     this.conf = conf;
     this.timeout = conf.getTimeout();
     this.channelFactory = channelFactory;
@@ -83,7 +90,8 @@ public abstract class AbstractGRPCClient<
                   return ClientCalls.blockingUnaryCall(
                       stub.getChannel(), method, stub.getCallOptions(), requestFactory.get());
                 },
-                method.getFullMethodName());
+                method.getFullMethodName(),
+                backOffer);
 
     if (logger.isTraceEnabled()) {
       logger.trace(String.format("leaving %s...", method.getFullMethodName()));
@@ -104,14 +112,15 @@ public abstract class AbstractGRPCClient<
         .create(handler)
         .callWithRetry(
             () -> {
-              StubT stub = getAsyncStub();
+              FutureStubT stub = getAsyncStub();
               ClientCalls.asyncUnaryCall(
                   stub.getChannel().newCall(method, stub.getCallOptions()),
                   requestFactory.get(),
                   responseObserver);
               return null;
             },
-            method.getFullMethodName());
+            method.getFullMethodName(),
+            backOffer);
     logger.debug(String.format("leaving %s...", method.getFullMethodName()));
   }
 
@@ -128,11 +137,12 @@ public abstract class AbstractGRPCClient<
             .create(handler)
             .callWithRetry(
                 () -> {
-                  StubT stub = getAsyncStub();
+                  FutureStubT stub = getAsyncStub();
                   return asyncBidiStreamingCall(
                       stub.getChannel().newCall(method, stub.getCallOptions()), responseObserver);
                 },
-                method.getFullMethodName());
+                method.getFullMethodName(),
+                backOffer);
     logger.debug(String.format("leaving %s...", method.getFullMethodName()));
     return observer;
   }
@@ -155,7 +165,8 @@ public abstract class AbstractGRPCClient<
                       blockingServerStreamingCall(
                           stub.getChannel(), method, stub.getCallOptions(), requestFactory.get()));
                 },
-                method.getFullMethodName());
+                method.getFullMethodName(),
+                backOffer);
     logger.debug(String.format("leaving %s...", method.getFullMethodName()));
     return response;
   }
@@ -170,5 +181,21 @@ public abstract class AbstractGRPCClient<
 
   protected abstract BlockingStubT getBlockingStub();
 
-  protected abstract StubT getAsyncStub();
+  protected abstract FutureStubT 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

@@ -15,14 +15,19 @@
 
 package org.tikv.common;
 
+import org.tikv.common.util.BackOffer;
 import org.tikv.kvproto.Kvrpcpb;
 
 public class ConfigUtils {
+  public static final String TIKV_CONFIGURATION_FILENAME = "tikv.properties";
+
   public static final String TIKV_PD_ADDRESSES = "tikv.pd.addresses";
   public static final String TIKV_GRPC_TIMEOUT = "tikv.grpc.timeout_in_ms";
+  public static final String TIKV_GRPC_FORWARD_TIMEOUT = "tikv.grpc.forward_timeout_in_ms";
   public static final String TIKV_GRPC_SCAN_TIMEOUT = "tikv.grpc.scan_timeout_in_ms";
   public static final String TIKV_GRPC_SCAN_BATCH_SIZE = "tikv.grpc.scan_batch_size";
   public static final String TIKV_GRPC_MAX_FRAME_SIZE = "tikv.grpc.max_frame_size";
+  public static final String TIKV_GRPC_IDLE_TIMEOUT = "tikv.grpc.idle_timeout";
 
   public static final String TIKV_INDEX_SCAN_BATCH_SIZE = "tikv.index.scan_batch_size";
   public static final String TIKV_INDEX_SCAN_CONCURRENCY = "tikv.index.scan_concurrency";
@@ -42,16 +47,54 @@ 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 TIKV_HEALTH_CHECK_PERIOD_DURATION =
+      "tikv.health_check_period_duration";
 
+  public static final String TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS = "tikv.rawkv.default_backoff_in_ms";
+  public static final String TIKV_RAWKV_READ_TIMEOUT_IN_MS = "tikv.rawkv.read_timeout_in_ms";
+  public static final String TIKV_RAWKV_WRITE_TIMEOUT_IN_MS = "tikv.rawkv.write_timeout_in_ms";
+  public static final String TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS =
+      "tikv.rawkv.batch_read_timeout_in_ms";
+  public static final String TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS =
+      "tikv.rawkv.batch_write_timeout_in_ms";
+  public static final String TIKV_RAWKV_SCAN_TIMEOUT_IN_MS = "tikv.rawkv.scan_timeout_in_ms";
+  public static final String TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS = "tikv.rawkv.clean_timeout_in_ms";
+  public static final String TIKV_BO_REGION_MISS_BASE_IN_MS = "tikv.bo_region_miss_base_in_ms";
+  public static final String TIKV_RAWKV_READ_SLOWLOG_IN_MS = "tikv.rawkv.read_slowlog_in_ms";
+  public static final String TIKV_RAWKV_WRITE_SLOWLOG_IN_MS = "tikv.rawkv.write_slowlog_in_ms";
+  public static final String TIKV_RAWKV_BATCH_READ_SLOWLOG_IN_MS =
+      "tikv.rawkv.batch_read_slowlog_in_ms";
+  public static final String TIKV_RAWKV_BATCH_WRITE_SLOWLOG_IN_MS =
+      "tikv.rawkv.batch_write_slowlog_in_ms";
+  public static final String TIKV_RAWKV_SCAN_SLOWLOG_IN_MS = "tikv.rawkv.scan_slowlog_in_ms";
+
+  public static final String TiKV_CIRCUIT_BREAK_ENABLE = "tikv.circuit_break.enable";
+  public static final String TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS =
+      "tikv.circuit_break.trigger.availability.window_in_seconds";
+  public static final String TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE =
+      "tikv.circuit_break.trigger.availability.error_threshold_percentage";
+  public static final String TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUEST_VOLUMN_THRESHOLD =
+      "tikv.circuit_break.trigger.availability.request_volumn_threshold";
+  public static final String TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS =
+      "tikv.circuit_break.trigger.sleep_window_in_seconds";
+  public static final String TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT =
+      "tikv.circuit_break.trigger.attempt_request_count";
+
+  public static final String TIFLASH_ENABLE = "tiflash.enable";
   public static final String DEF_PD_ADDRESSES = "127.0.0.1:2379";
-  public static final String DEF_TIMEOUT = "600ms";
+  public static final String DEF_TIMEOUT = "200ms";
+  public static final String DEF_FORWARD_TIMEOUT = "300ms";
   public static final String DEF_SCAN_TIMEOUT = "20s";
+  public static final int DEF_CHECK_HEALTH_TIMEOUT = 100;
+  public static final int DEF_HEALTH_CHECK_PERIOD_DURATION = 300;
   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 +115,23 @@ 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 int DEF_TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS = BackOffer.RAWKV_MAX_BACKOFF;
+
+  public static final int DEF_TIKV_RAWKV_READ_TIMEOUT_IN_MS = 2000;
+  public static final int DEF_TIKV_RAWKV_WRITE_TIMEOUT_IN_MS = 2000;
+  public static final int DEF_TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS = 2000;
+  public static final int DEF_TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS = 2000;
+  public static final int DEF_TIKV_RAWKV_SCAN_TIMEOUT_IN_MS = 10000;
+  public static final int DEF_TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS = 600000;
+
+  public static final int DEF_TIKV_BO_REGION_MISS_BASE_IN_MS = 20;
+  public static final String DEF_TIKV_RAWKV_SCAN_SLOWLOG_IN_MS = "5000";
 
   public static final String NORMAL_COMMAND_PRIORITY = "NORMAL";
   public static final String LOW_COMMAND_PRIORITY = "LOW";
@@ -86,4 +142,17 @@ 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";
+
+  public static final int DEF_TIKV_GRPC_IDLE_TIMEOUT = 60;
+
+  public static final boolean DEF_TiKV_CIRCUIT_BREAK_ENABLE = false;
+  public static final int DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS = 60;
+  public static final int DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE = 100;
+  public static final int DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUST_VOLUMN_THRESHOLD = 10;
+  public static final int DEF_TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS = 20;
+  public static final int DEF_TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT = 10;
 }

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

@@ -0,0 +1,87 @@
+/*
+ * 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;
+
+import static org.tikv.common.pd.PDUtils.addrToUri;
+
+import com.google.common.annotations.Beta;
+import io.etcd.jetcd.ByteSequence;
+import io.etcd.jetcd.Client;
+import io.etcd.jetcd.KeyValue;
+import io.etcd.jetcd.kv.GetResponse;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class DefaultHostMapping implements HostMapping {
+  private static final String NETWORK_MAPPING_PATH = "/client/url-mapping";
+  private final Client etcdClient;
+  private final String networkMappingName;
+  private final ConcurrentMap<String, String> hostMapping;
+  private final Logger logger = LoggerFactory.getLogger(DefaultHostMapping.class);
+
+  public DefaultHostMapping(Client etcdClient, String networkMappingName) {
+    this.etcdClient = etcdClient;
+    this.networkMappingName = networkMappingName;
+    this.hostMapping = new ConcurrentHashMap<>();
+  }
+
+  private ByteSequence hostToNetworkMappingKey(String host) {
+    String path = NETWORK_MAPPING_PATH + "/" + networkMappingName + "/" + host;
+    return ByteSequence.from(path, StandardCharsets.UTF_8);
+  }
+
+  @Beta
+  private String getMappedHostFromPD(String host) {
+    ByteSequence hostKey = hostToNetworkMappingKey(host);
+    for (int i = 0; i < 5; i++) {
+      CompletableFuture<GetResponse> future = etcdClient.getKVClient().get(hostKey);
+      try {
+        GetResponse resp = future.get();
+        List<KeyValue> kvs = resp.getKvs();
+        if (kvs.size() != 1) {
+          break;
+        }
+        return kvs.get(0).getValue().toString(StandardCharsets.UTF_8);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      } catch (ExecutionException e) {
+        logger.info("failed to get mapped Host from PD: " + host, e);
+        break;
+      } catch (Exception ignore) {
+        // ignore
+        break;
+      }
+    }
+    return host;
+  }
+
+  public URI getMappedURI(URI uri) {
+    if (networkMappingName.isEmpty()) {
+      return uri;
+    }
+    return addrToUri(
+        hostMapping.computeIfAbsent(uri.getHost(), this::getMappedHostFromPD)
+            + ":"
+            + uri.getPort());
+  }
+}

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

@@ -15,73 +15,9 @@
 
 package org.tikv.common;
 
-import static org.tikv.common.pd.PDUtils.addrToUri;
-
-import com.google.common.annotations.Beta;
-import io.etcd.jetcd.ByteSequence;
-import io.etcd.jetcd.Client;
-import io.etcd.jetcd.KeyValue;
-import io.etcd.jetcd.kv.GetResponse;
+import java.io.Serializable;
 import java.net.URI;
-import java.nio.charset.StandardCharsets;
-import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.ExecutionException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
-public class HostMapping {
-  private static final String NETWORK_MAPPING_PATH = "/client/url-mapping";
-  private final Client etcdClient;
-  private final String networkMappingName;
-  private final ConcurrentMap<String, String> hostMapping;
-  private final Logger logger = LoggerFactory.getLogger(HostMapping.class);
-
-  public HostMapping(Client etcdClient, String networkMappingName) {
-    this.etcdClient = etcdClient;
-    this.networkMappingName = networkMappingName;
-    this.hostMapping = new ConcurrentHashMap<>();
-  }
-
-  private ByteSequence hostToNetworkMappingKey(String host) {
-    String path = NETWORK_MAPPING_PATH + "/" + networkMappingName + "/" + host;
-    return ByteSequence.from(path, StandardCharsets.UTF_8);
-  }
-
-  @Beta
-  private String getMappedHostFromPD(String host) {
-    ByteSequence hostKey = hostToNetworkMappingKey(host);
-    for (int i = 0; i < 5; i++) {
-      CompletableFuture<GetResponse> future = etcdClient.getKVClient().get(hostKey);
-      try {
-        GetResponse resp = future.get();
-        List<KeyValue> kvs = resp.getKvs();
-        if (kvs.size() != 1) {
-          break;
-        }
-        return kvs.get(0).getValue().toString(StandardCharsets.UTF_8);
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-      } catch (ExecutionException e) {
-        logger.info("failed to get mapped Host from PD: " + host, e);
-        break;
-      } catch (Exception ignore) {
-        // ignore
-        break;
-      }
-    }
-    return host;
-  }
-
-  public URI getMappedURI(URI uri) {
-    if (networkMappingName.isEmpty()) {
-      return uri;
-    }
-    return addrToUri(
-        hostMapping.computeIfAbsent(uri.getHost(), this::getMappedHostFromPD)
-            + ":"
-            + uri.getPort());
-  }
+public interface HostMapping extends Serializable {
+  URI getMappedURI(URI uri);
 }

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

@@ -0,0 +1,80 @@
+/*
+ * 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;
+
+import io.prometheus.client.exporter.HTTPServer;
+import io.prometheus.client.hotspot.DefaultExports;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class MetricsServer {
+  private static final Logger logger = LoggerFactory.getLogger(MetricsServer.class);
+
+  private static MetricsServer METRICS_SERVER_INSTANCE = null;
+  private static int metricsServerRefCount = 0;
+
+  private int port;
+  private HTTPServer server;
+
+  public static MetricsServer getInstance(TiConfiguration conf) {
+    if (!conf.isMetricsEnable()) {
+      return null;
+    }
+
+    synchronized (MetricsServer.class) {
+      int port = conf.getMetricsPort();
+      if (METRICS_SERVER_INSTANCE != null) {
+        if (port != METRICS_SERVER_INSTANCE.port) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Do dot support multiple tikv.metrics.port, which are %d and %d",
+                  port, METRICS_SERVER_INSTANCE.port));
+        }
+      } else {
+        METRICS_SERVER_INSTANCE = new MetricsServer(port);
+      }
+      metricsServerRefCount += 1;
+      return METRICS_SERVER_INSTANCE;
+    }
+  }
+
+  private MetricsServer(int port) {
+    try {
+      this.port = port;
+      DefaultExports.initialize();
+      this.server = new HTTPServer(port, true);
+      logger.info("http server is up " + this.server.getPort());
+    } catch (Exception e) {
+      logger.error("http server not up");
+      throw new RuntimeException(e);
+    }
+  }
+
+  public void close() {
+    synchronized (MetricsServer.class) {
+      if (metricsServerRefCount == 1) {
+        if (server != null) {
+          server.stop();
+          logger.info("Metrics server on " + server.getPort() + " is stopped");
+        }
+        METRICS_SERVER_INSTANCE = null;
+      }
+
+      if (metricsServerRefCount >= 1) {
+        metricsServerRefCount -= 1;
+      }
+    }
+  }
+}

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

@@ -30,23 +30,29 @@ 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;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 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;
 import org.tikv.common.codec.Codec.BytesCodec;
+import org.tikv.common.codec.CodecDataInput;
 import org.tikv.common.codec.CodecDataOutput;
 import org.tikv.common.codec.KeyUtils;
 import org.tikv.common.exception.GrpcException;
@@ -54,16 +60,17 @@ import org.tikv.common.exception.TiClientInternalException;
 import org.tikv.common.meta.TiTimestamp;
 import org.tikv.common.operation.NoopHandler;
 import org.tikv.common.operation.PDErrorHandler;
-import org.tikv.common.region.TiRegion;
 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.common.util.Pair;
+import org.tikv.kvproto.Metapb;
 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.PDGrpc.PDFutureStub;
+import org.tikv.kvproto.Pdpb;
 import org.tikv.kvproto.Pdpb.Error;
 import org.tikv.kvproto.Pdpb.ErrorType;
 import org.tikv.kvproto.Pdpb.GetAllStoresRequest;
@@ -85,19 +92,21 @@ import org.tikv.kvproto.Pdpb.Timestamp;
 import org.tikv.kvproto.Pdpb.TsoRequest;
 import org.tikv.kvproto.Pdpb.TsoResponse;
 
-public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
+public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDFutureStub>
     implements ReadOnlyPDClient {
   private static final String TIFLASH_TABLE_SYNC_PROGRESS_PATH = "/tiflash/table/sync";
+  private static final long MIN_TRY_UPDATE_DURATION = 50;
   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;
   private Client etcdClient;
   private ConcurrentMap<Long, Double> tiflashReplicaMap;
   private HostMapping hostMapping;
+  private long lastUpdateLeaderTime;
 
   public static final Histogram PD_GET_REGION_BY_KEY_REQUEST_LATENCY =
       Histogram.build()
@@ -143,7 +152,7 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
    *
    * @param region represents a region info
    */
-  void scatterRegion(TiRegion region, BackOffer backOffer) {
+  void scatterRegion(Metapb.Region region, BackOffer backOffer) {
     Supplier<ScatterRegionRequest> request =
         () ->
             ScatterRegionRequest.newBuilder().setHeader(header).setRegionId(region.getId()).build();
@@ -167,7 +176,7 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
    *
    * @param region
    */
-  void waitScatterRegionFinish(TiRegion region, BackOffer backOffer) {
+  void waitScatterRegionFinish(Metapb.Region region, BackOffer backOffer) {
     for (; ; ) {
       GetOperatorResponse resp = getOperator(region.getId());
       if (resp != null) {
@@ -220,7 +229,7 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   }
 
   @Override
-  public TiRegion getRegionByKey(BackOffer backOffer, ByteString key) {
+  public Pair<Metapb.Region, Metapb.Peer> getRegionByKey(BackOffer backOffer, ByteString key) {
     Histogram.Timer requestTimer = PD_GET_REGION_BY_KEY_REQUEST_LATENCY.startTimer();
     try {
       if (conf.getKvMode() == KVMode.TXN) {
@@ -238,42 +247,14 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
 
       GetRegionResponse resp =
           callWithRetry(backOffer, PDGrpc.getGetRegionMethod(), request, handler);
-      return new TiRegion(
-          resp.getRegion(),
-          resp.getLeader(),
-          conf.getIsolationLevel(),
-          conf.getCommandPriority(),
-          conf.getKvMode(),
-          conf.isReplicaRead());
+      return new Pair<Metapb.Region, Metapb.Peer>(decodeRegion(resp.getRegion()), resp.getLeader());
     } 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) {
+  public Pair<Metapb.Region, Metapb.Peer> getRegionByID(BackOffer backOffer, long id) {
     Supplier<GetRegionByIDRequest> request =
         () -> GetRegionByIDRequest.newBuilder().setHeader(header).setRegionId(id).build();
     PDErrorHandler<GetRegionResponse> handler =
@@ -281,37 +262,7 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
 
     GetRegionResponse resp =
         callWithRetry(backOffer, PDGrpc.getGetRegionByIDMethod(), request, handler);
-    // Instead of using default leader instance, explicitly set no leader to null
-    return new TiRegion(
-        resp.getRegion(),
-        resp.getLeader(),
-        conf.getIsolationLevel(),
-        conf.getCommandPriority(),
-        conf.getKvMode(),
-        conf.isReplicaRead());
-  }
-
-  @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();
+    return new Pair<Metapb.Region, Metapb.Peer>(decodeRegion(resp.getRegion()), resp.getLeader());
   }
 
   private Supplier<GetStoreRequest> buildGetStoreReq(long storeId) {
@@ -329,23 +280,16 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
 
   @Override
   public Store getStore(BackOffer backOffer, long storeId) {
-    return callWithRetry(
-            backOffer, PDGrpc.getGetStoreMethod(), buildGetStoreReq(storeId), buildPDErrorHandler())
-        .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();
+    GetStoreResponse resp =
+        callWithRetry(
+            backOffer,
+            PDGrpc.getGetStoreMethod(),
+            buildGetStoreReq(storeId),
+            buildPDErrorHandler());
+    if (resp != null) {
+      return resp.getStore();
+    }
+    return null;
   }
 
   @Override
@@ -361,8 +305,8 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   }
 
   @Override
-  public boolean isReplicaRead() {
-    return conf.isReplicaRead();
+  public TiConfiguration.ReplicaRead getReplicaRead() {
+    return conf.getReplicaRead();
   }
 
   @Override
@@ -385,14 +329,15 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
   }
 
   @VisibleForTesting
-  LeaderWrapper getLeaderWrapper() {
-    return leaderWrapper;
+  PDClientWrapper getPdClientWrapper() {
+    return pdClientWrapper;
   }
 
   private GetMembersResponse getMembers(URI uri) {
     try {
       ManagedChannel probChan = channelFactory.getChannel(uriToAddr(uri), hostMapping);
-      PDGrpc.PDBlockingStub stub = PDGrpc.newBlockingStub(probChan);
+      PDGrpc.PDBlockingStub stub =
+          PDGrpc.newBlockingStub(probChan).withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
       GetMembersRequest request =
           GetMembersRequest.newBuilder().setHeader(RequestHeader.getDefaultInstance()).build();
       GetMembersResponse resp = stub.getMembers(request);
@@ -407,55 +352,147 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
     return null;
   }
 
-  synchronized boolean switchLeader(List<String> leaderURLs) {
-    if (leaderURLs.isEmpty()) return false;
-    String leaderUrlStr = leaderURLs.get(0);
-    // TODO: Why not strip protocol info on server side since grpc does not need it
-    if (leaderWrapper != null && leaderUrlStr.equals(leaderWrapper.getLeaderInfo())) {
-      return true;
+  // return whether the leader has changed to target address `leaderUrlStr`.
+  synchronized boolean trySwitchLeader(String leaderUrlStr) {
+    if (pdClientWrapper != null) {
+      if (leaderUrlStr.equals(pdClientWrapper.getLeaderInfo())) {
+        // The message to leader is not forwarded by follower.
+        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 =
-          new LeaderWrapper(
-              leaderUrlStr,
-              PDGrpc.newBlockingStub(clientChannel),
-              PDGrpc.newStub(clientChannel),
-              System.nanoTime());
+      pdClientWrapper =
+          new PDClientWrapper(leaderUrlStr, leaderUrlStr, clientChannel, System.nanoTime());
+      timeout = conf.getTimeout();
     } 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());
+      timeout = conf.getForwardTimeout();
+    } catch (IllegalArgumentException e) {
+      return false;
+    }
+    logger.info(String.format("Switched to new leader by follower forward: %s", pdClientWrapper));
+    return true;
+  }
+
+  public synchronized void updateLeaderOrforwardFollower() {
+    if (System.currentTimeMillis() - lastUpdateLeaderTime < MIN_TRY_UPDATE_DURATION) {
+      return;
+    }
     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)) {
+        lastUpdateLeaderTime = System.currentTimeMillis();
+        return;
+      }
+
+      if (!conf.getEnableGrpcForward()) {
+        continue;
+      }
+
+      logger.info(String.format("can not switch to new leader, try follower forward"));
+      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)) {
+          logger.warn(
+              String.format("forward request to pd [%s] by pd [%s]", leaderUrlStr, followerUrlStr));
+          return;
+        }
+      }
+    }
+    lastUpdateLeaderTime = System.currentTimeMillis();
+    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;
       }
     }
-    throw new TiClientInternalException(
-        "already tried all address on file, but not leader found yet.");
+    lastUpdateLeaderTime = System.currentTimeMillis();
+    if (pdClientWrapper == null) {
+      throw new TiClientInternalException(
+          "already tried all address on file, but not leader found yet.");
+    }
+  }
+
+  private synchronized void tryUpdateMembers(List<URI> members) {
+    this.pdAddrs = members;
   }
 
   public void updateTiFlashReplicaStatus() {
@@ -513,74 +550,125 @@ 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) {
+  protected PDFutureStub getAsyncStub() {
+    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() {
+    logger.info("init cluster: start");
     GetMembersResponse resp = null;
-    List<URI> pdAddrs = getConf().getPdAddrs();
+    List<URI> pdAddrs = new ArrayList<>(getConf().getPdAddrs());
+    // shuffle PD addresses so that clients call getMembers from different PD
+    Collections.shuffle(pdAddrs);
     this.pdAddrs = pdAddrs;
-    this.etcdClient = Client.builder().endpoints(pdAddrs).build();
-    this.hostMapping = new HostMapping(this.etcdClient, conf.getNetworkMappingName());
+    this.etcdClient =
+        Client.builder()
+            .endpoints(pdAddrs)
+            .executorService(
+                Executors.newCachedThreadPool(
+                    new ThreadFactoryBuilder()
+                        .setNameFormat("etcd-conn-manager-pool-%d")
+                        .setDaemon(true)
+                        .build()))
+            .build();
+    logger.info("init host mapping: start");
+    this.hostMapping =
+        Optional.ofNullable(getConf().getHostMapping())
+            .orElseGet(() -> new DefaultHostMapping(this.etcdClient, conf.getNetworkMappingName()));
+    logger.info("init host mapping: end");
+    // The first request may cost too much latency
+    long originTimeout = this.timeout;
+    this.timeout = 2000;
     for (URI u : pdAddrs) {
+      logger.info("get members with pd " + u + ": start");
       resp = getMembers(u);
+      logger.info("get members with pd " + u + ": end");
       if (resp != null) {
         break;
       }
     }
+    if (resp == null) {
+      logger.error("Could not get leader member with: " + pdAddrs);
+    }
+
+    this.timeout = originTimeout;
     checkNotNull(resp, "Failed to init client for PD cluster.");
     long clusterId = resp.getHeader().getClusterId();
     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));
+    logger.info("createLeaderClientWrapper with leader " + leaderUrlStr + ": start");
+    createLeaderClientWrapper(leaderUrlStr);
+    logger.info("createLeaderClientWrapper with leader " + leaderUrlStr + ": end");
     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,
-        1,
-        TimeUnit.MINUTES);
+        10,
+        10,
+        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);
+    logger.info("init cluster: finish");
   }
 
-  static class LeaderWrapper {
+  static class PDClientWrapper {
     private final String leaderInfo;
     private final PDBlockingStub blockingStub;
-    private final PDStub asyncStub;
+    private final PDFutureStub asyncStub;
     private final long createTime;
+    private final String storeAddress;
 
-    LeaderWrapper(
-        String leaderInfo,
-        PDGrpc.PDBlockingStub blockingStub,
-        PDGrpc.PDStub asyncStub,
-        long createTime) {
+    PDClientWrapper(
+        String leaderInfo, String storeAddress, ManagedChannel clientChannel, long createTime) {
+      if (!storeAddress.equals(leaderInfo)) {
+        Metadata header = new Metadata();
+        header.put(TiConfiguration.PD_FORWARD_META_DATA_KEY, addrToUri(leaderInfo).toString());
+        this.blockingStub =
+            MetadataUtils.attachHeaders(PDGrpc.newBlockingStub(clientChannel), header);
+        this.asyncStub = MetadataUtils.attachHeaders(PDGrpc.newFutureStub(clientChannel), header);
+      } else {
+        this.blockingStub = PDGrpc.newBlockingStub(clientChannel);
+        this.asyncStub = PDGrpc.newFutureStub(clientChannel);
+      }
       this.leaderInfo = leaderInfo;
-      this.blockingStub = blockingStub;
-      this.asyncStub = asyncStub;
+      this.storeAddress = storeAddress;
       this.createTime = createTime;
     }
 
@@ -588,11 +676,15 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
       return leaderInfo;
     }
 
+    String getStoreAddress() {
+      return storeAddress;
+    }
+
     PDBlockingStub getBlockingStub() {
       return blockingStub;
     }
 
-    PDStub getAsyncStub() {
+    PDFutureStub getAsyncStub() {
       return asyncStub;
     }
 
@@ -602,7 +694,32 @@ public class PDClient extends AbstractGRPCClient<PDBlockingStub, PDStub>
 
     @Override
     public String toString() {
-      return "[leaderInfo: " + leaderInfo + "]";
+      return "[leaderInfo: " + leaderInfo + ", storeAddress: " + storeAddress + "]";
     }
   }
+
+  private Metapb.Region decodeRegion(Metapb.Region region) {
+    final boolean isRawRegion = conf.getKvMode() == KVMode.RAW;
+    Metapb.Region.Builder builder =
+        Metapb.Region.newBuilder()
+            .setId(region.getId())
+            .setRegionEpoch(region.getRegionEpoch())
+            .addAllPeers(region.getPeersList());
+
+    if (region.getStartKey().isEmpty() || isRawRegion) {
+      builder.setStartKey(region.getStartKey());
+    } else {
+      byte[] decodedStartKey = BytesCodec.readBytes(new CodecDataInput(region.getStartKey()));
+      builder.setStartKey(ByteString.copyFrom(decodedStartKey));
+    }
+
+    if (region.getEndKey().isEmpty() || isRawRegion) {
+      builder.setEndKey(region.getEndKey());
+    } else {
+      byte[] decodedEndKey = BytesCodec.readBytes(new CodecDataInput(region.getEndKey()));
+      builder.setEndKey(ByteString.copyFrom(decodedEndKey));
+    }
+
+    return builder.build();
+  }
 }

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

@@ -17,10 +17,10 @@ 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;
+import org.tikv.common.util.Pair;
+import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.Metapb.Store;
 
 /** Readonly PD client including only reading related interface Supposed for TiDB-like use cases */
@@ -38,9 +38,7 @@ public interface ReadOnlyPDClient {
    * @param key key in bytes for locating a region
    * @return the region whose startKey and endKey range covers the given key
    */
-  TiRegion getRegionByKey(BackOffer backOffer, ByteString key);
-
-  Future<TiRegion> getRegionByKeyAsync(BackOffer backOffer, ByteString key);
+  Pair<Metapb.Region, Metapb.Peer> getRegionByKey(BackOffer backOffer, ByteString key);
 
   /**
    * Get Region by Region Id
@@ -48,9 +46,7 @@ public interface ReadOnlyPDClient {
    * @param id Region Id
    * @return the region corresponding to the given Id
    */
-  TiRegion getRegionByID(BackOffer backOffer, long id);
-
-  Future<TiRegion> getRegionByIDAsync(BackOffer backOffer, long id);
+  Pair<Metapb.Region, Metapb.Peer> getRegionByID(BackOffer backOffer, long id);
 
   HostMapping getHostMapping();
 
@@ -62,9 +58,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,9 @@ package org.tikv.common;
 
 import static org.tikv.common.ConfigUtils.*;
 
+import io.grpc.Metadata;
+import java.io.IOException;
+import java.io.InputStream;
 import java.io.Serializable;
 import java.net.URI;
 import java.util.*;
@@ -24,6 +27,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,10 +35,17 @@ 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<String> FORWARD_META_DATA_KEY =
+      Metadata.Key.of("tikv-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
+  public static final Metadata.Key<String> PD_FORWARD_META_DATA_KEY =
+      Metadata.Key.of("pd-forwarded-host", Metadata.ASCII_STRING_MARSHALLER);
 
   static {
+    // priority: system environment > config file > default
     loadFromSystemProperties();
+    loadFromConfigurationFile();
     loadFromDefaultProperties();
+    listAll();
   }
 
   private static void loadFromSystemProperties() {
@@ -45,9 +56,35 @@ public class TiConfiguration implements Serializable {
     }
   }
 
+  private static void loadFromConfigurationFile() {
+    try (InputStream input =
+        TiConfiguration.class
+            .getClassLoader()
+            .getResourceAsStream(ConfigUtils.TIKV_CONFIGURATION_FILENAME)) {
+      Properties properties = new Properties();
+
+      if (input == null) {
+        logger.warn("Unable to find " + ConfigUtils.TIKV_CONFIGURATION_FILENAME);
+        return;
+      }
+
+      logger.info("loading " + ConfigUtils.TIKV_CONFIGURATION_FILENAME);
+      properties.load(input);
+      for (String key : properties.stringPropertyNames()) {
+        if (key.startsWith("tikv.")) {
+          String value = properties.getProperty(key);
+          setIfMissing(key, value);
+        }
+      }
+    } catch (IOException e) {
+      logger.error("load config file error", e);
+    }
+  }
+
   private static void loadFromDefaultProperties() {
     setIfMissing(TIKV_PD_ADDRESSES, DEF_PD_ADDRESSES);
     setIfMissing(TIKV_GRPC_TIMEOUT, DEF_TIMEOUT);
+    setIfMissing(TIKV_GRPC_FORWARD_TIMEOUT, DEF_FORWARD_TIMEOUT);
     setIfMissing(TIKV_GRPC_SCAN_TIMEOUT, DEF_SCAN_TIMEOUT);
     setIfMissing(TIKV_GRPC_SCAN_BATCH_SIZE, DEF_SCAN_BATCH_SIZE);
     setIfMissing(TIKV_GRPC_MAX_FRAME_SIZE, DEF_MAX_FRAME_SIZE);
@@ -67,14 +104,41 @@ 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);
+    setIfMissing(TIKV_HEALTH_CHECK_PERIOD_DURATION, DEF_HEALTH_CHECK_PERIOD_DURATION);
+    setIfMissing(TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS, DEF_TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS);
+    setIfMissing(TIKV_GRPC_IDLE_TIMEOUT, DEF_TIKV_GRPC_IDLE_TIMEOUT);
+    setIfMissing(TIKV_RAWKV_READ_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_READ_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_RAWKV_WRITE_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_WRITE_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_RAWKV_SCAN_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_SCAN_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS, DEF_TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS);
+    setIfMissing(TIKV_BO_REGION_MISS_BASE_IN_MS, DEF_TIKV_BO_REGION_MISS_BASE_IN_MS);
+    setIfMissing(TIKV_RAWKV_SCAN_SLOWLOG_IN_MS, DEF_TIKV_RAWKV_SCAN_SLOWLOG_IN_MS);
+    setIfMissing(TiKV_CIRCUIT_BREAK_ENABLE, DEF_TiKV_CIRCUIT_BREAK_ENABLE);
+    setIfMissing(
+        TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS,
+        DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS);
+    setIfMissing(
+        TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE,
+        DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE);
+    setIfMissing(
+        TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUEST_VOLUMN_THRESHOLD,
+        DEF_TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUST_VOLUMN_THRESHOLD);
+    setIfMissing(
+        TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS, DEF_TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS);
+    setIfMissing(
+        TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT, DEF_TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT);
   }
 
   public static void listAll() {
-    logger.info(new ArrayList<>(settings.entrySet()).toString());
+    logger.info("static configurations are:" + new ArrayList<>(settings.entrySet()).toString());
   }
 
   private static void set(String key, String value) {
@@ -117,10 +181,14 @@ public class TiConfiguration implements Serializable {
     return option.get();
   }
 
-  private static int getInt(String key) {
+  public static int getInt(String key) {
     return Integer.parseInt(get(key));
   }
 
+  public static Optional<Integer> getIntOption(String key) {
+    return getOption(key).map(Integer::parseInt);
+  }
+
   private static int getInt(String key, int defaultValue) {
     try {
       return getOption(key).map(Integer::parseInt).orElse(defaultValue);
@@ -216,7 +284,19 @@ 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 forwardTimeout = getTimeAsMs(TIKV_GRPC_FORWARD_TIMEOUT);
   private long scanTimeout = getTimeAsMs(TIKV_GRPC_SCAN_TIMEOUT);
   private int maxFrameSize = getInt(TIKV_GRPC_MAX_FRAME_SIZE);
   private List<URI> pdAddrs = getPdAddrs(TIKV_PD_ADDRESSES);
@@ -233,20 +313,57 @@ 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 int healthCheckPeriodDuration = getInt(TIKV_HEALTH_CHECK_PERIOD_DURATION);
 
   private final String networkMappingName = get(TIKV_NETWORK_MAPPING_NAME);
+  private HostMapping hostMapping = null;
+
+  private int rawKVDefaultBackoffInMS = getInt(TIKV_RAWKV_DEFAULT_BACKOFF_IN_MS);
+  private int rawKVReadTimeoutInMS = getInt(TIKV_RAWKV_READ_TIMEOUT_IN_MS);
+  private int rawKVWriteTimeoutInMS = getInt(TIKV_RAWKV_WRITE_TIMEOUT_IN_MS);
+  private int rawKVBatchReadTimeoutInMS = getInt(TIKV_RAWKV_BATCH_READ_TIMEOUT_IN_MS);
+  private int rawKVBatchWriteTimeoutInMS = getInt(TIKV_RAWKV_BATCH_WRITE_TIMEOUT_IN_MS);
+  private int rawKVScanTimeoutInMS = getInt(TIKV_RAWKV_SCAN_TIMEOUT_IN_MS);
+  private int rawKVCleanTimeoutInMS = getInt(TIKV_RAWKV_CLEAN_TIMEOUT_IN_MS);
+  private Optional<Integer> rawKVReadSlowLogInMS = getIntOption(TIKV_RAWKV_READ_SLOWLOG_IN_MS);
+  private Optional<Integer> rawKVWriteSlowLogInMS = getIntOption(TIKV_RAWKV_WRITE_SLOWLOG_IN_MS);
+  private Optional<Integer> rawKVBatchReadSlowLogInMS =
+      getIntOption(TIKV_RAWKV_BATCH_READ_SLOWLOG_IN_MS);
+  private Optional<Integer> rawKVBatchWriteSlowLogInMS =
+      getIntOption(TIKV_RAWKV_BATCH_WRITE_SLOWLOG_IN_MS);
+  private int rawKVScanSlowLogInMS = getInt(TIKV_RAWKV_SCAN_SLOWLOG_IN_MS);
+  private int idleTimeout = getInt(TIKV_GRPC_IDLE_TIMEOUT);
+  private boolean circuitBreakEnable = getBoolean(TiKV_CIRCUIT_BREAK_ENABLE);
+  private int circuitBreakAvailabilityWindowInSeconds =
+      getInt(TiKV_CIRCUIT_BREAK_AVAILABILITY_WINDOW_IN_SECONDS);
+  private int circuitBreakAvailabilityErrorThresholdPercentage =
+      getInt(TiKV_CIRCUIT_BREAK_AVAILABILITY_ERROR_THRESHOLD_PERCENTAGE);
+  private int circuitBreakAvailabilityRequestVolumnThreshold =
+      getInt(TiKV_CIRCUIT_BREAK_AVAILABILITY_REQUEST_VOLUMN_THRESHOLD);
+  private int circuitBreakSleepWindowInSeconds = getInt(TiKV_CIRCUIT_BREAK_SLEEP_WINDOW_IN_SECONDS);
+  private int circuitBreakAttemptRequestCount = getInt(TiKV_CIRCUIT_BREAK_ATTEMPT_REQUEST_COUNT);
 
   public enum KVMode {
     TXN,
     RAW
   }
 
+  public enum ReplicaRead {
+    LEADER,
+    FOLLOWER,
+    LEADER_AND_FOLLOWER
+  }
+
   public static TiConfiguration createDefault() {
     return new TiConfiguration();
   }
@@ -301,6 +418,15 @@ public class TiConfiguration implements Serializable {
     return this;
   }
 
+  public long getForwardTimeout() {
+    return forwardTimeout;
+  }
+
+  public TiConfiguration setForwardTimeout(long timeout) {
+    this.forwardTimeout = timeout;
+    return this;
+  }
+
   public long getScanTimeout() {
     return scanTimeout;
   }
@@ -457,15 +583,40 @@ public class TiConfiguration implements Serializable {
     return this;
   }
 
-  public boolean isReplicaRead() {
-    return isReplicaRead;
+  public ReplicaRead getReplicaRead() {
+    return replicaRead;
   }
 
-  public TiConfiguration setReplicaRead(boolean isReplicaRead) {
-    this.isReplicaRead = isReplicaRead;
+  public TiConfiguration setReplicaRead(ReplicaRead replicaRead) {
+    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 +638,181 @@ public class TiConfiguration implements Serializable {
   public String getNetworkMappingName() {
     return this.networkMappingName;
   }
+
+  public HostMapping getHostMapping() {
+    return hostMapping;
+  }
+
+  public void setHostMapping(HostMapping mapping) {
+    this.hostMapping = mapping;
+  }
+
+  public boolean getEnableGrpcForward() {
+    return this.enableGrpcForward;
+  }
+
+  public long getGrpcHealthCheckTimeout() {
+    return this.grpcHealthCheckTimeout;
+  }
+
+  public long getHealthCheckPeriodDuration() {
+    return this.healthCheckPeriodDuration;
+  }
+
+  public int getRawKVDefaultBackoffInMS() {
+    return rawKVDefaultBackoffInMS;
+  }
+
+  public void setRawKVDefaultBackoffInMS(int rawKVDefaultBackoffInMS) {
+    this.rawKVDefaultBackoffInMS = rawKVDefaultBackoffInMS;
+  }
+
+  public int getIdleTimeout() {
+    return idleTimeout;
+  }
+
+  public void setIdleTimeout(int timeout) {
+    this.idleTimeout = timeout;
+  }
+
+  public int getRawKVReadTimeoutInMS() {
+    return rawKVReadTimeoutInMS;
+  }
+
+  public void setRawKVReadTimeoutInMS(int rawKVReadTimeoutInMS) {
+    this.rawKVReadTimeoutInMS = rawKVReadTimeoutInMS;
+  }
+
+  public int getRawKVWriteTimeoutInMS() {
+    return rawKVWriteTimeoutInMS;
+  }
+
+  public void setRawKVWriteTimeoutInMS(int rawKVWriteTimeoutInMS) {
+    this.rawKVWriteTimeoutInMS = rawKVWriteTimeoutInMS;
+  }
+
+  public int getRawKVBatchReadTimeoutInMS() {
+    return rawKVBatchReadTimeoutInMS;
+  }
+
+  public void setRawKVBatchReadTimeoutInMS(int rawKVBatchReadTimeoutInMS) {
+    this.rawKVBatchReadTimeoutInMS = rawKVBatchReadTimeoutInMS;
+  }
+
+  public int getRawKVBatchWriteTimeoutInMS() {
+    return rawKVBatchWriteTimeoutInMS;
+  }
+
+  public void setRawKVBatchWriteTimeoutInMS(int rawKVBatchWriteTimeoutInMS) {
+    this.rawKVBatchWriteTimeoutInMS = rawKVBatchWriteTimeoutInMS;
+  }
+
+  public int getRawKVScanTimeoutInMS() {
+    return rawKVScanTimeoutInMS;
+  }
+
+  public void setRawKVScanTimeoutInMS(int rawKVScanTimeoutInMS) {
+    this.rawKVScanTimeoutInMS = rawKVScanTimeoutInMS;
+  }
+
+  public int getRawKVCleanTimeoutInMS() {
+    return rawKVCleanTimeoutInMS;
+  }
+
+  public void setRawKVCleanTimeoutInMS(int rawKVCleanTimeoutInMS) {
+    this.rawKVCleanTimeoutInMS = rawKVCleanTimeoutInMS;
+  }
+
+  public Integer getRawKVReadSlowLogInMS() {
+    return rawKVReadSlowLogInMS.orElse((int) (getTimeout() * 2));
+  }
+
+  public void setRawKVReadSlowLogInMS(Integer rawKVReadSlowLogInMS) {
+    this.rawKVReadSlowLogInMS = Optional.of(rawKVReadSlowLogInMS);
+  }
+
+  public Integer getRawKVWriteSlowLogInMS() {
+    return rawKVWriteSlowLogInMS.orElse((int) (getTimeout() * 2));
+  }
+
+  public void setRawKVWriteSlowLogInMS(Integer rawKVWriteSlowLogInMS) {
+    this.rawKVWriteSlowLogInMS = Optional.of(rawKVWriteSlowLogInMS);
+  }
+
+  public Integer getRawKVBatchReadSlowLogInMS() {
+    return rawKVBatchReadSlowLogInMS.orElse((int) (getTimeout() * 2));
+  }
+
+  public void setRawKVBatchReadSlowLogInMS(Integer rawKVBatchReadSlowLogInMS) {
+    this.rawKVBatchReadSlowLogInMS = Optional.of(rawKVBatchReadSlowLogInMS);
+  }
+
+  public Integer getRawKVBatchWriteSlowLogInMS() {
+    return rawKVBatchWriteSlowLogInMS.orElse((int) (getTimeout() * 2));
+  }
+
+  public void setRawKVBatchWriteSlowLogInMS(Integer rawKVBatchWriteSlowLogInMS) {
+    this.rawKVBatchWriteSlowLogInMS = Optional.of(rawKVBatchWriteSlowLogInMS);
+  }
+
+  public int getRawKVScanSlowLogInMS() {
+    return rawKVScanSlowLogInMS;
+  }
+
+  public void setRawKVScanSlowLogInMS(int rawKVScanSlowLogInMS) {
+    this.rawKVScanSlowLogInMS = rawKVScanSlowLogInMS;
+  }
+
+  public boolean isCircuitBreakEnable() {
+    return circuitBreakEnable;
+  }
+
+  public void setCircuitBreakEnable(boolean circuitBreakEnable) {
+    this.circuitBreakEnable = circuitBreakEnable;
+  }
+
+  public int getCircuitBreakAvailabilityWindowInSeconds() {
+    return circuitBreakAvailabilityWindowInSeconds;
+  }
+
+  public void setCircuitBreakAvailabilityWindowInSeconds(
+      int circuitBreakAvailabilityWindowInSeconds) {
+    this.circuitBreakAvailabilityWindowInSeconds = circuitBreakAvailabilityWindowInSeconds;
+  }
+
+  public int getCircuitBreakAvailabilityErrorThresholdPercentage() {
+    return circuitBreakAvailabilityErrorThresholdPercentage;
+  }
+
+  public void setCircuitBreakAvailabilityErrorThresholdPercentage(
+      int circuitBreakAvailabilityErrorThresholdPercentage) {
+    this.circuitBreakAvailabilityErrorThresholdPercentage =
+        circuitBreakAvailabilityErrorThresholdPercentage;
+  }
+
+  public int getCircuitBreakAvailabilityRequestVolumnThreshold() {
+    return circuitBreakAvailabilityRequestVolumnThreshold;
+  }
+
+  public void setCircuitBreakAvailabilityRequestVolumnThreshold(
+      int circuitBreakAvailabilityRequestVolumnThreshold) {
+    this.circuitBreakAvailabilityRequestVolumnThreshold =
+        circuitBreakAvailabilityRequestVolumnThreshold;
+  }
+
+  public int getCircuitBreakSleepWindowInSeconds() {
+    return circuitBreakSleepWindowInSeconds;
+  }
+
+  public void setCircuitBreakSleepWindowInSeconds(int circuitBreakSleepWindowInSeconds) {
+    this.circuitBreakSleepWindowInSeconds = circuitBreakSleepWindowInSeconds;
+  }
+
+  public int getCircuitBreakAttemptRequestCount() {
+    return circuitBreakAttemptRequestCount;
+  }
+
+  public void setCircuitBreakAttemptRequestCount(int circuitBreakAttemptRequestCount) {
+    this.circuitBreakAttemptRequestCount = circuitBreakAttemptRequestCount;
+  }
 }

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

@@ -20,32 +20,30 @@ import static org.tikv.common.util.ClientUtils.groupKeysByRegion;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.protobuf.ByteString;
-import io.prometheus.client.CollectorRegistry;
-import io.prometheus.client.exporter.HTTPServer;
-import java.net.InetSocketAddress;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.catalog.Catalog;
-import org.tikv.common.event.CacheInvalidateEvent;
 import org.tikv.common.exception.TiKVException;
 import org.tikv.common.key.Key;
 import org.tikv.common.meta.TiTimestamp;
-import org.tikv.common.policy.RetryPolicy;
 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.raw.SmartRawKVClient;
+import org.tikv.service.failsafe.CircuitBreaker;
+import org.tikv.service.failsafe.CircuitBreakerImpl;
 import org.tikv.txn.KVClient;
 import org.tikv.txn.TxnKVClient;
 
@@ -59,7 +57,6 @@ public class TiSession implements AutoCloseable {
   private static final Map<String, TiSession> sessionCachedMap = new HashMap<>();
   private final TiConfiguration conf;
   private final ChannelFactory channelFactory;
-  private Function<CacheInvalidateEvent, Void> cacheInvalidateCallback;
   // below object creation is either heavy or making connection (pd), pending for lazy loading
   private volatile PDClient client;
   private volatile Catalog catalog;
@@ -71,37 +68,63 @@ 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;
-  private CollectorRegistry collectorRegistry;
+  private volatile boolean isClosed = false;
+  private MetricsServer metricsServer;
+  private final CircuitBreaker circuitBreaker;
 
   public TiSession(TiConfiguration conf) {
+    // may throw org.tikv.common.MetricsServer  - http server not up
+    // put it at the beginning of this function to avoid unclosed Thread
+    this.metricsServer = MetricsServer.getInstance(conf);
+
     this.conf = conf;
-    this.channelFactory = new ChannelFactory(conf.getMaxFrameSize());
+    this.channelFactory = new ChannelFactory(conf.getMaxFrameSize(), conf.getIdleTimeout());
     this.client = PDClient.createRaw(conf, channelFactory);
-    if (conf.isMetricsEnable()) {
-      try {
-        this.collectorRegistry = new CollectorRegistry();
-        this.collectorRegistry.register(RawKVClient.RAW_REQUEST_LATENCY);
-        this.collectorRegistry.register(RawKVClient.RAW_REQUEST_FAILURE);
-        this.collectorRegistry.register(RawKVClient.RAW_REQUEST_SUCCESS);
-        this.collectorRegistry.register(RegionStoreClient.GRPC_RAW_REQUEST_LATENCY);
-        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.server =
-            new HTTPServer(
-                new InetSocketAddress(conf.getMetricsPort()), this.collectorRegistry, true);
-        logger.info("http server is up " + this.server.getPort());
-      } catch (Exception e) {
-        logger.error("http server not up");
-        throw new RuntimeException(e);
-      }
+    this.enableGrpcForward = conf.getEnableGrpcForward();
+    if (this.enableGrpcForward) {
+      logger.info("enable grpc forward for high available");
     }
+    warmUp();
+    this.circuitBreaker = new CircuitBreakerImpl(conf);
     logger.info("TiSession initialized in " + conf.getKvMode() + " mode");
   }
 
+  private synchronized void warmUp() {
+    long warmUpStartTime = System.currentTimeMillis();
+    try {
+      this.client = getPDClient();
+      this.regionManager = getRegionManager();
+      List<Metapb.Store> stores = this.client.getAllStores(ConcreteBackOffer.newGetBackOff());
+      // warm up store cache
+      for (Metapb.Store store : stores) {
+        this.regionManager.updateStore(
+            null,
+            new TiStore(this.client.getStore(ConcreteBackOffer.newGetBackOff(), store.getId())));
+      }
+      ByteString startKey = ByteString.EMPTY;
+
+      do {
+        TiRegion region = regionManager.getRegionByKey(startKey);
+        startKey = region.getEndKey();
+      } while (!startKey.isEmpty());
+
+      RawKVClient rawKVClient = createRawClient();
+      ByteString exampleKey = ByteString.EMPTY;
+      ByteString prev = rawKVClient.get(exampleKey);
+      rawKVClient.delete(exampleKey);
+      rawKVClient.putIfAbsent(exampleKey, prev);
+      rawKVClient.put(exampleKey, prev);
+    } catch (Exception e) {
+      // ignore error
+      logger.info("warm up fails, ignored ", e);
+    } finally {
+      logger.info(
+          String.format("warm up duration %d ms", System.currentTimeMillis() - warmUpStartTime));
+    }
+  }
+
   @VisibleForTesting
   public static TiSession create(TiConfiguration conf) {
     return new TiSession(conf);
@@ -122,22 +145,35 @@ public class TiSession implements AutoCloseable {
   }
 
   public RawKVClient createRawClient() {
+    checkIsClosed();
+
     RegionStoreClientBuilder builder =
         new RegionStoreClientBuilder(conf, channelFactory, this.getRegionManager(), client);
     return new RawKVClient(this, builder);
   }
 
+  public SmartRawKVClient createSmartRawClient() {
+    RawKVClient rawKVClient = createRawClient();
+    return new SmartRawKVClient(rawKVClient, circuitBreaker);
+  }
+
   public KVClient createKVClient() {
+    checkIsClosed();
+
     RegionStoreClientBuilder builder =
         new RegionStoreClientBuilder(conf, channelFactory, this.getRegionManager(), client);
     return new KVClient(conf, builder);
   }
 
   public TxnKVClient createTxnClient() {
+    checkIsClosed();
+
     return new TxnKVClient(conf, this.getRegionStoreClientBuilder(), this.getPDClient());
   }
 
   public RegionStoreClient.RegionStoreClientBuilder getRegionStoreClientBuilder() {
+    checkIsClosed();
+
     RegionStoreClient.RegionStoreClientBuilder res = clientBuilder;
     if (res == null) {
       synchronized (this) {
@@ -157,18 +193,26 @@ public class TiSession implements AutoCloseable {
   }
 
   public TiTimestamp getTimestamp() {
+    checkIsClosed();
+
     return getPDClient().getTimestamp(ConcreteBackOffer.newTsoBackOff());
   }
 
   public Snapshot createSnapshot() {
+    checkIsClosed();
+
     return new Snapshot(getTimestamp(), this);
   }
 
   public Snapshot createSnapshot(TiTimestamp ts) {
+    checkIsClosed();
+
     return new Snapshot(ts, this);
   }
 
   public PDClient getPDClient() {
+    checkIsClosed();
+
     PDClient res = client;
     if (res == null) {
       synchronized (this) {
@@ -182,6 +226,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public Catalog getCatalog() {
+    checkIsClosed();
+
     Catalog res = catalog;
     if (res == null) {
       synchronized (this) {
@@ -194,12 +240,14 @@ public class TiSession implements AutoCloseable {
     return res;
   }
 
-  public synchronized RegionManager getRegionManager() {
+  public RegionManager getRegionManager() {
+    checkIsClosed();
+
     RegionManager res = regionManager;
     if (res == null) {
       synchronized (this) {
         if (regionManager == null) {
-          regionManager = new RegionManager(getPDClient(), this.cacheInvalidateCallback);
+          regionManager = new RegionManager(getConf(), getPDClient(), this.channelFactory);
         }
         res = regionManager;
       }
@@ -208,6 +256,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForIndexScan() {
+    checkIsClosed();
+
     ExecutorService res = indexScanThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -227,6 +277,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForTableScan() {
+    checkIsClosed();
+
     ExecutorService res = tableScanThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -243,6 +295,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForBatchPut() {
+    checkIsClosed();
+
     ExecutorService res = batchPutThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -262,6 +316,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForBatchGet() {
+    checkIsClosed();
+
     ExecutorService res = batchGetThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -281,6 +337,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForBatchDelete() {
+    checkIsClosed();
+
     ExecutorService res = batchDeleteThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -300,6 +358,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForBatchScan() {
+    checkIsClosed();
+
     ExecutorService res = batchScanThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -319,6 +379,8 @@ public class TiSession implements AutoCloseable {
   }
 
   public ExecutorService getThreadPoolForDeleteRange() {
+    checkIsClosed();
+
     ExecutorService res = deleteRangeThreadPool;
     if (res == null) {
       synchronized (this) {
@@ -339,22 +401,11 @@ public class TiSession implements AutoCloseable {
 
   @VisibleForTesting
   public ChannelFactory getChannelFactory() {
+    checkIsClosed();
+
     return channelFactory;
   }
 
-  public CollectorRegistry getCollectorRegistry() {
-    return collectorRegistry;
-  }
-
-  /**
-   * This is used for setting call back function to invalidate cache information
-   *
-   * @param callBackFunc callback function
-   */
-  public void injectCallBackFunc(Function<CacheInvalidateEvent, Void> callBackFunc) {
-    this.cacheInvalidateCallback = callBackFunc;
-  }
-
   /**
    * split region and scatter
    *
@@ -365,11 +416,13 @@ public class TiSession implements AutoCloseable {
       int splitRegionBackoffMS,
       int scatterRegionBackoffMS,
       int scatterWaitMS) {
+    checkIsClosed();
+
     logger.info(String.format("split key's size is %d", splitKeys.size()));
     long startMS = System.currentTimeMillis();
 
     // split region
-    List<TiRegion> newRegions =
+    List<Metapb.Region> newRegions =
         splitRegion(
             splitKeys
                 .stream()
@@ -378,7 +431,7 @@ public class TiSession implements AutoCloseable {
             ConcreteBackOffer.newCustomBackOff(splitRegionBackoffMS));
 
     // scatter region
-    for (TiRegion newRegion : newRegions) {
+    for (Metapb.Region newRegion : newRegions) {
       try {
         getPDClient()
             .scatterRegion(newRegion, ConcreteBackOffer.newCustomBackOff(scatterRegionBackoffMS));
@@ -391,7 +444,7 @@ public class TiSession implements AutoCloseable {
     if (scatterWaitMS > 0) {
       logger.info("start to wait scatter region finish");
       long scatterRegionStartMS = System.currentTimeMillis();
-      for (TiRegion newRegion : newRegions) {
+      for (Metapb.Region newRegion : newRegions) {
         long remainMS = (scatterRegionStartMS + scatterWaitMS) - System.currentTimeMillis();
         if (remainMS <= 0) {
           logger.warn("wait scatter region timeout");
@@ -408,17 +461,17 @@ public class TiSession implements AutoCloseable {
     logger.info("splitRegionAndScatter cost {} seconds", (endMS - startMS) / 1000);
   }
 
-  private List<TiRegion> splitRegion(List<ByteString> splitKeys, BackOffer backOffer) {
-    List<TiRegion> regions = new ArrayList<>();
+  private List<Metapb.Region> splitRegion(List<ByteString> splitKeys, BackOffer backOffer) {
+    List<Metapb.Region> regions = new ArrayList<>();
 
     Map<TiRegion, List<ByteString>> groupKeys =
         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()
@@ -431,7 +484,7 @@ public class TiSession implements AutoCloseable {
             "split key equal to region start key or end key. Region splitting is not needed.");
       } else {
         logger.info("start to split region id={}, split size={}", region.getId(), splits.size());
-        List<TiRegion> newRegions;
+        List<Metapb.Region> newRegions;
         try {
           newRegions = getRegionStoreClientBuilder().build(region, store).splitRegion(splits);
         } catch (final TiKVException e) {
@@ -450,49 +503,115 @@ public class TiSession implements AutoCloseable {
     return regions;
   }
 
-  @Override
-  public synchronized void close() throws Exception {
+  private void checkIsClosed() {
     if (isClosed) {
-      logger.warn("this TiSession is already closed!");
-      return;
-    }
-
-    if (server != null) {
-      server.stop();
-      logger.info("Metrics server on " + server.getPort() + " is stopped");
-    }
-
-    isClosed = true;
-    synchronized (sessionCachedMap) {
-      sessionCachedMap.remove(conf.getPdAddrsString());
-    }
-
-    if (tableScanThreadPool != null) {
-      tableScanThreadPool.shutdownNow();
-    }
-    if (indexScanThreadPool != null) {
-      indexScanThreadPool.shutdownNow();
-    }
-    if (batchGetThreadPool != null) {
-      batchGetThreadPool.shutdownNow();
-    }
-    if (batchPutThreadPool != null) {
-      batchPutThreadPool.shutdownNow();
-    }
-    if (batchDeleteThreadPool != null) {
-      batchDeleteThreadPool.shutdownNow();
-    }
-    if (batchScanThreadPool != null) {
-      batchScanThreadPool.shutdownNow();
-    }
-    if (deleteRangeThreadPool != null) {
-      deleteRangeThreadPool.shutdownNow();
-    }
-    if (client != null) {
-      getPDClient().close();
-    }
-    if (catalog != null) {
-      getCatalog().close();
+      throw new RuntimeException("this TiSession is closed!");
     }
   }
+
+  public synchronized void closeAwaitTermination(long timeoutMS) throws Exception {
+    shutdown(false);
+
+    long startMS = System.currentTimeMillis();
+    while (true) {
+      if (isTerminatedExecutorServices()) {
+        cleanAfterTerminated();
+        return;
+      }
+
+      if (System.currentTimeMillis() - startMS > timeoutMS) {
+        shutdown(true);
+        return;
+      }
+      Thread.sleep(500);
+    }
+  }
+
+  @Override
+  public synchronized void close() throws Exception {
+    shutdown(true);
+  }
+
+  private synchronized void shutdown(boolean now) throws Exception {
+    if (!isClosed) {
+      isClosed = true;
+      synchronized (sessionCachedMap) {
+        sessionCachedMap.remove(conf.getPdAddrsString());
+      }
+
+      if (metricsServer != null) {
+        metricsServer.close();
+      }
+
+      if (circuitBreaker != null) {
+        circuitBreaker.close();
+      }
+    }
+
+    if (now) {
+      shutdownNowExecutorServices();
+      cleanAfterTerminated();
+    } else {
+      shutdownExecutorServices();
+    }
+  }
+
+  private synchronized void cleanAfterTerminated() throws InterruptedException {
+    if (regionManager != null) {
+      regionManager.close();
+    }
+    if (client != null) {
+      client.close();
+    }
+    if (catalog != null) {
+      catalog.close();
+    }
+  }
+
+  private List<ExecutorService> getExecutorServices() {
+    List<ExecutorService> executorServiceList = new ArrayList<>();
+    if (tableScanThreadPool != null) {
+      executorServiceList.add(tableScanThreadPool);
+    }
+    if (indexScanThreadPool != null) {
+      executorServiceList.add(indexScanThreadPool);
+    }
+    if (batchGetThreadPool != null) {
+      executorServiceList.add(batchGetThreadPool);
+    }
+    if (batchPutThreadPool != null) {
+      executorServiceList.add(batchPutThreadPool);
+    }
+    if (batchDeleteThreadPool != null) {
+      executorServiceList.add(batchDeleteThreadPool);
+    }
+    if (batchScanThreadPool != null) {
+      executorServiceList.add(batchScanThreadPool);
+    }
+    if (deleteRangeThreadPool != null) {
+      executorServiceList.add(deleteRangeThreadPool);
+    }
+    return executorServiceList;
+  }
+
+  private void shutdownExecutorServices() {
+    for (ExecutorService executorService : getExecutorServices()) {
+      executorService.shutdown();
+    }
+  }
+
+  private void shutdownNowExecutorServices() {
+    for (ExecutorService executorService : getExecutorServices()) {
+      executorService.shutdownNow();
+    }
+  }
+
+  private boolean isTerminatedExecutorServices() {
+    for (ExecutorService executorService : getExecutorServices()) {
+      if (!executorService.isTerminated()) {
+        return false;
+      }
+    }
+    return true;
+  }
 }

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/CircuitBreakerOpenException.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2018 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.exception;
+
+public class CircuitBreakerOpenException extends RuntimeException {
+  public CircuitBreakerOpenException() {
+    super("Circuit Breaker Opened");
+  }
+}

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

@@ -0,0 +1,23 @@
+/*
+ * 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.exception;
+
+public class InvalidStoreException extends TiKVException {
+
+  public InvalidStoreException(long storeId) {
+    super(String.format("Invalid storeId: %d", storeId));
+  }
+}

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/log/SlowLog.java

@@ -0,0 +1,28 @@
+/*
+ *
+ * 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.log;
+
+public interface SlowLog {
+  void addProperty(String key, String value);
+
+  SlowLogSpan start(String name);
+
+  void setError(Throwable err);
+
+  void log();
+}

src/main/java/org/tikv/common/log/SlowLogEmptyImpl.java

@@ -0,0 +1,38 @@
+/*
+ *
+ * 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.log;
+
+public class SlowLogEmptyImpl implements SlowLog {
+  public static final SlowLogEmptyImpl INSTANCE = new SlowLogEmptyImpl();
+
+  private SlowLogEmptyImpl() {}
+
+  @Override
+  public void addProperty(String key, String value) {}
+
+  @Override
+  public SlowLogSpan start(String name) {
+    return SlowLogSpanEmptyImpl.INSTANCE;
+  }
+
+  @Override
+  public void setError(Throwable err) {}
+
+  @Override
+  public void log() {}
+}

src/main/java/org/tikv/common/log/SlowLogImpl.java

@@ -0,0 +1,102 @@
+/*
+ *
+ * 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.log;
+
+import com.google.gson.JsonArray;
+import com.google.gson.JsonObject;
+import java.text.SimpleDateFormat;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class SlowLogImpl implements SlowLog {
+  private static final Logger logger = LoggerFactory.getLogger(SlowLogImpl.class);
+
+  private static final int MAX_SPAN_SIZE = 1024;
+
+  public static final SimpleDateFormat DATE_FORMAT = new SimpleDateFormat("HH:mm:ss.SSS");
+
+  private final List<SlowLogSpan> slowLogSpans = new ArrayList<>();
+  private Throwable error = null;
+
+  private final long startMS;
+  private final long slowThresholdMS;
+
+  /** Key-Value pairs which will be logged, e.g. function name, key, region, etc. */
+  private final Map<String, String> properties;
+
+  public SlowLogImpl(long slowThresholdMS, Map<String, String> properties) {
+    this.startMS = System.currentTimeMillis();
+    this.slowThresholdMS = slowThresholdMS;
+    this.properties = new HashMap<>(properties);
+  }
+
+  @Override
+  public void addProperty(String key, String value) {
+    this.properties.put(key, value);
+  }
+
+  @Override
+  public synchronized SlowLogSpan start(String name) {
+    SlowLogSpan slowLogSpan = new SlowLogSpanImpl(name);
+    if (slowLogSpans.size() < MAX_SPAN_SIZE) {
+      slowLogSpans.add(slowLogSpan);
+    }
+    slowLogSpan.start();
+    return slowLogSpan;
+  }
+
+  @Override
+  public void setError(Throwable err) {
+    this.error = err;
+  }
+
+  @Override
+  public void log() {
+    long currentMS = System.currentTimeMillis();
+    if (error != null || (slowThresholdMS >= 0 && currentMS - startMS > slowThresholdMS)) {
+      logger.warn("SlowLog:" + getSlowLogString(currentMS));
+    }
+  }
+
+  private String getSlowLogString(long currentMS) {
+    JsonObject jsonObject = new JsonObject();
+
+    jsonObject.addProperty("start", DATE_FORMAT.format(startMS));
+    jsonObject.addProperty("end", DATE_FORMAT.format(currentMS));
+    jsonObject.addProperty("duration", (currentMS - startMS) + "ms");
+    if (error != null) {
+      jsonObject.addProperty("error", error.getMessage());
+    }
+
+    for (Map.Entry<String, String> entry : properties.entrySet()) {
+      jsonObject.addProperty(entry.getKey(), entry.getValue());
+    }
+
+    JsonArray jsonArray = new JsonArray();
+    for (SlowLogSpan slowLogSpan : slowLogSpans) {
+      jsonArray.add(slowLogSpan.toJsonElement());
+    }
+    jsonObject.add("spans", jsonArray);
+
+    return jsonObject.toString();
+  }
+}

src/main/java/org/tikv/common/log/SlowLogSpan.java

@@ -0,0 +1,28 @@
+/*
+ *
+ * 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.log;
+
+import com.google.gson.JsonElement;
+
+public interface SlowLogSpan {
+  void start();
+
+  void end();
+
+  JsonElement toJsonElement();
+}

src/main/java/org/tikv/common/log/SlowLogSpanEmptyImpl.java

@@ -0,0 +1,39 @@
+/*
+ *
+ * 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.log;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+
+public class SlowLogSpanEmptyImpl implements SlowLogSpan {
+
+  public static final SlowLogSpanEmptyImpl INSTANCE = new SlowLogSpanEmptyImpl();
+
+  private SlowLogSpanEmptyImpl() {}
+
+  @Override
+  public void start() {}
+
+  @Override
+  public void end() {}
+
+  @Override
+  public JsonElement toJsonElement() {
+    return new JsonObject();
+  }
+}

src/main/java/org/tikv/common/log/SlowLogSpanImpl.java

@@ -0,0 +1,77 @@
+/*
+ *
+ * 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.log;
+
+import static org.tikv.common.log.SlowLogImpl.DATE_FORMAT;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+
+public class SlowLogSpanImpl implements SlowLogSpan {
+  private final String name;
+  private long startMS;
+  private long endMS;
+
+  public SlowLogSpanImpl(String name) {
+    this.name = name;
+    this.startMS = 0;
+    this.endMS = 0;
+  }
+
+  @Override
+  public void start() {
+    this.startMS = System.currentTimeMillis();
+  }
+
+  @Override
+  public void end() {
+    this.endMS = System.currentTimeMillis();
+  }
+
+  @Override
+  public JsonElement toJsonElement() {
+    JsonObject jsonObject = new JsonObject();
+    jsonObject.addProperty("name", name);
+    jsonObject.addProperty("start", getStartString());
+    jsonObject.addProperty("end", getEndString());
+    jsonObject.addProperty("duration", getDurationString());
+
+    return jsonObject;
+  }
+
+  private String getStartString() {
+    if (startMS == 0) {
+      return "N/A";
+    }
+    return DATE_FORMAT.format(startMS);
+  }
+
+  private String getEndString() {
+    if (endMS == 0) {
+      return "N/A";
+    }
+    return DATE_FORMAT.format(endMS);
+  }
+
+  private String getDurationString() {
+    if (startMS == 0 || endMS == 0) {
+      return "N/A";
+    }
+    return (endMS - startMS) + "ms";
+  }
+}

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

@@ -19,20 +19,14 @@ package org.tikv.common.operation;
 
 import static org.tikv.common.util.BackOffFunction.BackOffFuncType.BoTxnLockFast;
 
-import com.google.protobuf.ByteString;
-import io.grpc.Status;
-import io.grpc.StatusRuntimeException;
 import java.util.Collections;
 import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.tikv.common.codec.KeyUtils;
 import org.tikv.common.exception.GrpcException;
 import org.tikv.common.exception.KeyException;
 import org.tikv.common.region.RegionErrorReceiver;
 import org.tikv.common.region.RegionManager;
-import org.tikv.common.region.TiRegion;
-import org.tikv.common.util.BackOffFunction;
 import org.tikv.common.util.BackOffer;
 import org.tikv.kvproto.Errorpb;
 import org.tikv.kvproto.Kvrpcpb;
@@ -43,16 +37,12 @@ import org.tikv.txn.ResolveLockResult;
 // TODO: consider refactor to Builder mode
 public class KVErrorHandler<RespT> implements ErrorHandler<RespT> {
   private static final Logger logger = LoggerFactory.getLogger(KVErrorHandler.class);
-  // if a store does not have leader currently, store id is set to 0
-  private static final int NO_LEADER_STORE_ID = 0;
-  private final Function<RespT, Errorpb.Error> getRegionError;
   private final Function<RespT, Kvrpcpb.KeyError> getKeyError;
   private final Function<ResolveLockResult, Object> resolveLockResultCallback;
-  private final RegionManager regionManager;
-  private final RegionErrorReceiver recv;
   private final AbstractLockResolverClient lockResolverClient;
   private final long callerStartTS;
   private final boolean forWrite;
+  private final RegionErrorHandler<RespT> regionHandler;
 
   public KVErrorHandler(
       RegionManager regionManager,
@@ -63,42 +53,14 @@ public class KVErrorHandler<RespT> implements ErrorHandler<RespT> {
       Function<ResolveLockResult, Object> resolveLockResultCallback,
       long callerStartTS,
       boolean forWrite) {
-    this.recv = recv;
+    this.regionHandler = new RegionErrorHandler<>(regionManager, recv, getRegionError);
     this.lockResolverClient = lockResolverClient;
-    this.regionManager = regionManager;
-    this.getRegionError = getRegionError;
     this.getKeyError = getKeyError;
     this.resolveLockResultCallback = resolveLockResultCallback;
     this.callerStartTS = callerStartTS;
     this.forWrite = forWrite;
   }
 
-  public KVErrorHandler(
-      RegionManager regionManager,
-      RegionErrorReceiver recv,
-      Function<RespT, Errorpb.Error> getRegionError) {
-    this.recv = recv;
-    this.lockResolverClient = null;
-    this.regionManager = regionManager;
-    this.getRegionError = getRegionError;
-    this.getKeyError = resp -> null;
-    this.resolveLockResultCallback = resolveLock -> null;
-    this.callerStartTS = 0;
-    this.forWrite = false;
-  }
-
-  private Errorpb.Error getRegionError(RespT resp) {
-    if (getRegionError != null) {
-      return getRegionError.apply(resp);
-    }
-    return null;
-  }
-
-  private void invalidateRegionStoreCache(TiRegion ctxRegion) {
-    regionManager.invalidateRegion(ctxRegion);
-    regionManager.invalidateStore(ctxRegion.getLeader().getStoreId());
-  }
-
   private void resolveLock(BackOffer backOffer, Lock lock) {
     if (lockResolverClient != null) {
       logger.warn("resolving lock");
@@ -124,160 +86,32 @@ public class KVErrorHandler<RespT> implements ErrorHandler<RespT> {
   public boolean handleResponseError(BackOffer backOffer, RespT resp) {
     if (resp == null) {
       String msg =
-          String.format(
-              "Request Failed with unknown reason for region region [%s]", recv.getRegion());
+          String.format("Request Failed with unknown reason for [%s]", regionHandler.getRegion());
       logger.warn(msg);
       return handleRequestError(backOffer, new GrpcException(msg));
     }
 
-    // Region error handling logic
-    Errorpb.Error error = getRegionError(resp);
+    Errorpb.Error error = regionHandler.getRegionError(resp);
     if (error != null) {
-      if (error.hasNotLeader()) {
-        // this error is reported from raftstore:
-        // peer of current request is not leader, the following might be its causes:
-        // 1. cache is outdated, region has changed its leader, can be solved by re-fetching from PD
-        // 2. leader of current region is missing, need to wait and then fetch region info from PD
-        long newStoreId = error.getNotLeader().getLeader().getStoreId();
-        boolean retry;
-
-        // update Leader here
-        logger.warn(
-            String.format(
-                "NotLeader Error with region id %d and store id %d, new store id %d",
-                recv.getRegion().getId(), recv.getRegion().getLeader().getStoreId(), newStoreId));
-
-        BackOffFunction.BackOffFuncType backOffFuncType;
-        // if there's current no leader, we do not trigger update pd cache logic
-        // since issuing store = NO_LEADER_STORE_ID requests to pd will definitely fail.
-        if (newStoreId != NO_LEADER_STORE_ID) {
-          // If update leader fails, we need to fetch new region info from pd,
-          // and re-split key range for new region. Setting retry to false will
-          // stop retry and enter handleCopResponse logic, which would use RegionMiss
-          // backOff strategy to wait, fetch new region and re-split key range.
-          // onNotLeader is only needed when updateLeader succeeds, thus switch
-          // to a new store address.
-          TiRegion newRegion = this.regionManager.updateLeader(recv.getRegion(), newStoreId);
-          retry =
-              newRegion != null
-                  && recv.onNotLeader(this.regionManager.getStoreById(newStoreId), newRegion);
-
-          backOffFuncType = BackOffFunction.BackOffFuncType.BoUpdateLeader;
-        } else {
-          logger.info(
-              String.format(
-                  "Received zero store id, from region %d try next time",
-                  recv.getRegion().getId()));
-
-          backOffFuncType = BackOffFunction.BackOffFuncType.BoRegionMiss;
-          retry = false;
-        }
-
-        if (!retry) {
-          this.regionManager.invalidateRegion(recv.getRegion());
-        }
-
-        backOffer.doBackOff(backOffFuncType, new GrpcException(error.toString()));
-
-        return retry;
-      } else if (error.hasStoreNotMatch()) {
-        // this error is reported from raftstore:
-        // store_id requested at the moment is inconsistent with that expected
-        // Solution：re-fetch from PD
-        long storeId = recv.getRegion().getLeader().getStoreId();
-        long actualStoreId = error.getStoreNotMatch().getActualStoreId();
-        logger.warn(
-            String.format(
-                "Store Not Match happened with region id %d, store id %d, actual store id %d",
-                recv.getRegion().getId(), storeId, actualStoreId));
-
-        this.regionManager.invalidateRegion(recv.getRegion());
-        this.regionManager.invalidateStore(storeId);
-        // recv.onStoreNotMatch(this.regionManager.getStoreById(storeId));
-        // assume this is a low probability error, do not retry, just re-split the request by
-        // throwing it out.
-        return false;
-      } else if (error.hasEpochNotMatch()) {
-        // this error is reported from raftstore:
-        // region has outdated version，please try later.
-        logger.warn(String.format("Stale Epoch encountered for region [%s]", recv.getRegion()));
-        this.regionManager.onRegionStale(recv.getRegion());
-        return false;
-      } else if (error.hasServerIsBusy()) {
-        // this error is reported from kv:
-        // will occur when write pressure is high. Please try later.
-        logger.warn(
-            String.format(
-                "Server is busy for region [%s], reason: %s",
-                recv.getRegion(), error.getServerIsBusy().getReason()));
-        backOffer.doBackOff(
-            BackOffFunction.BackOffFuncType.BoServerBusy,
-            new StatusRuntimeException(
-                Status.fromCode(Status.Code.UNAVAILABLE).withDescription(error.toString())));
-        backOffer.doBackOff(
-            BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
-        return true;
-      } else if (error.hasStaleCommand()) {
-        // this error is reported from raftstore:
-        // command outdated, please try later
-        logger.warn(String.format("Stale command for region [%s]", recv.getRegion()));
-        backOffer.doBackOff(
-            BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
-        return true;
-      } else if (error.hasRaftEntryTooLarge()) {
-        logger.warn(String.format("Raft too large for region [%s]", recv.getRegion()));
-        throw new StatusRuntimeException(
-            Status.fromCode(Status.Code.UNAVAILABLE).withDescription(error.toString()));
-      } else if (error.hasKeyNotInRegion()) {
-        // this error is reported from raftstore:
-        // key requested is not in current region
-        // should not happen here.
-        ByteString invalidKey = error.getKeyNotInRegion().getKey();
-        logger.error(
-            String.format(
-                "Key not in region [%s] for key [%s], this error should not happen here.",
-                recv.getRegion(), KeyUtils.formatBytesUTF8(invalidKey)));
-        throw new StatusRuntimeException(Status.UNKNOWN.withDescription(error.toString()));
-      }
-
-      logger.warn(String.format("Unknown error %s for region [%s]", error, recv.getRegion()));
-      // For other errors, we only drop cache here.
-      // Upper level may split this task.
-      invalidateRegionStoreCache(recv.getRegion());
-      // retry if raft proposal is dropped, it indicates the store is in the middle of transition
-      if (error.getMessage().contains("Raft ProposalDropped")) {
-        backOffer.doBackOff(
-            BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
-        return true;
-      }
+      return regionHandler.handleRegionError(backOffer, error);
     }
 
-    boolean retry = false;
-
     // Key error handling logic
     Kvrpcpb.KeyError keyError = getKeyError.apply(resp);
     if (keyError != null) {
       try {
         Lock lock = AbstractLockResolverClient.extractLockFromKeyErr(keyError);
         resolveLock(backOffer, lock);
-        retry = true;
+        return true;
       } catch (KeyException e) {
         logger.warn("Unable to handle KeyExceptions other than LockException", e);
       }
     }
-    return retry;
+    return false;
   }
 
   @Override
   public boolean handleRequestError(BackOffer backOffer, Exception e) {
-    regionManager.onRequestFail(recv.getRegion());
-
-    backOffer.doBackOff(
-        BackOffFunction.BackOffFuncType.BoTiKVRPC,
-        new GrpcException(
-            "send tikv request error: " + e.getMessage() + ", try next peer later", e));
-    // TiKV maybe down, so do not retry in `callWithRetry`
-    // should re-fetch the new leader from PD and send request to it
-    return false;
+    return regionHandler.handleRequestError(backOffer, e);
   }
 }

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

@@ -19,6 +19,7 @@ package org.tikv.common.operation;
 
 import static org.tikv.common.pd.PDError.buildFromPdpbError;
 
+import io.grpc.StatusRuntimeException;
 import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -48,7 +49,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 +59,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());
@@ -72,7 +75,12 @@ public class PDErrorHandler<RespT> implements ErrorHandler<RespT> {
 
   @Override
   public boolean handleRequestError(BackOffer backOffer, Exception e) {
+    // store id is not found
+    if (e instanceof StatusRuntimeException && e.getMessage().contains("invalid store ID")) {
+      return false;
+    }
     backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoPDRPC, e);
+    client.updateLeaderOrforwardFollower();
     return true;
   }
 }

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

@@ -0,0 +1,254 @@
+package org.tikv.common.operation;
+
+import com.google.protobuf.ByteString;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.codec.KeyUtils;
+import org.tikv.common.exception.GrpcException;
+import org.tikv.common.exception.TiKVException;
+import org.tikv.common.region.RegionErrorReceiver;
+import org.tikv.common.region.RegionManager;
+import org.tikv.common.region.TiRegion;
+import org.tikv.common.util.BackOffFunction;
+import org.tikv.common.util.BackOffer;
+import org.tikv.kvproto.Errorpb;
+import org.tikv.kvproto.Metapb;
+
+public class RegionErrorHandler<RespT> implements ErrorHandler<RespT> {
+  private static final Logger logger = LoggerFactory.getLogger(RegionErrorHandler.class);
+  // if a store does not have leader currently, store id is set to 0
+  private static final int NO_LEADER_STORE_ID = 0;
+  private final Function<RespT, Errorpb.Error> getRegionError;
+  private final RegionManager regionManager;
+  private final RegionErrorReceiver recv;
+
+  public RegionErrorHandler(
+      RegionManager regionManager,
+      RegionErrorReceiver recv,
+      Function<RespT, Errorpb.Error> getRegionError) {
+    this.recv = recv;
+    this.regionManager = regionManager;
+    this.getRegionError = getRegionError;
+  }
+
+  @Override
+  public boolean handleResponseError(BackOffer backOffer, RespT resp) {
+    if (resp == null) {
+      String msg = String.format("Request Failed with unknown reason for [%s]", recv.getRegion());
+      return handleRequestError(backOffer, new GrpcException(msg));
+    }
+    // Region error handling logic
+    Errorpb.Error error = getRegionError(resp);
+    if (error != null) {
+      return handleRegionError(backOffer, error);
+    }
+    return false;
+  }
+
+  public boolean handleRegionError(BackOffer backOffer, Errorpb.Error error) {
+    if (error.hasNotLeader()) {
+      // this error is reported from raftstore:
+      // peer of current request is not leader, the following might be its causes:
+      // 1. cache is outdated, region has changed its leader, can be solved by re-fetching from PD
+      // 2. leader of current region is missing, need to wait and then fetch region info from PD
+      long newStoreId = error.getNotLeader().getLeader().getStoreId();
+      boolean retry;
+
+      // update Leader here
+      logger.warn(
+          String.format(
+              "NotLeader Error with region id %d and store id %d, new store id %d",
+              recv.getRegion().getId(), recv.getRegion().getLeader().getStoreId(), newStoreId));
+
+      BackOffFunction.BackOffFuncType backOffFuncType;
+      // if there's current no leader, we do not trigger update pd cache logic
+      // since issuing store = NO_LEADER_STORE_ID requests to pd will definitely fail.
+      if (newStoreId != NO_LEADER_STORE_ID) {
+        // If update leader fails, we need to fetch new region info from pd,
+        // and re-split key range for new region. Setting retry to false will
+        // stop retry and enter handleCopResponse logic, which would use RegionMiss
+        // backOff strategy to wait, fetch new region and re-split key range.
+        // onNotLeader is only needed when updateLeader succeeds, thus switch
+        // to a new store address.
+        TiRegion newRegion = this.regionManager.updateLeader(recv.getRegion(), newStoreId);
+        retry = newRegion != null && recv.onNotLeader(newRegion);
+
+        backOffFuncType = BackOffFunction.BackOffFuncType.BoUpdateLeader;
+      } else {
+        logger.info(
+            String.format(
+                "Received zero store id, from region %d try next time", recv.getRegion().getId()));
+
+        backOffFuncType = BackOffFunction.BackOffFuncType.BoRegionMiss;
+        retry = false;
+      }
+
+      if (!retry) {
+        this.regionManager.invalidateRegion(recv.getRegion());
+      }
+
+      backOffer.doBackOff(backOffFuncType, new GrpcException(error.toString()));
+
+      return retry;
+    } else if (error.hasStoreNotMatch()) {
+      // this error is reported from raftstore:
+      // store_id requested at the moment is inconsistent with that expected
+      // Solution：re-fetch from PD
+      long storeId = recv.getRegion().getLeader().getStoreId();
+      long actualStoreId = error.getStoreNotMatch().getActualStoreId();
+      logger.warn(
+          String.format(
+              "Store Not Match happened with region id %d, store id %d, actual store id %d",
+              recv.getRegion().getId(), storeId, actualStoreId));
+
+      this.regionManager.invalidateRegion(recv.getRegion());
+      this.regionManager.invalidateStore(storeId);
+      // recv.onStoreNotMatch(this.regionManager.getStoreById(storeId));
+      // assume this is a low probability error, do not retry, just re-split the request by
+      // throwing it out.
+      return false;
+    } else if (error.hasEpochNotMatch()) {
+      logger.warn(
+          String.format("tikv reports `EpochNotMatch` retry later, region: %s", recv.getRegion()));
+      return onRegionEpochNotMatch(backOffer, error.getEpochNotMatch().getCurrentRegionsList());
+    } else if (error.hasServerIsBusy()) {
+      // this error is reported from kv:
+      // will occur when write pressure is high. Please try later.
+      logger.warn(
+          String.format(
+              "Server is busy for region [%s], reason: %s",
+              recv.getRegion(), error.getServerIsBusy().getReason()));
+      backOffer.doBackOff(
+          BackOffFunction.BackOffFuncType.BoServerBusy,
+          new StatusRuntimeException(
+              Status.fromCode(Status.Code.UNAVAILABLE).withDescription(error.toString())));
+      backOffer.doBackOff(
+          BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
+      return true;
+    } else if (error.hasStaleCommand()) {
+      // this error is reported from raftstore:
+      // command outdated, please try later
+      logger.warn(String.format("Stale command for region [%s]", recv.getRegion()));
+      backOffer.doBackOff(
+          BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
+      return true;
+    } else if (error.hasRaftEntryTooLarge()) {
+      logger.warn(String.format("Raft too large for region [%s]", recv.getRegion()));
+      throw new StatusRuntimeException(
+          Status.fromCode(Status.Code.UNAVAILABLE).withDescription(error.toString()));
+    } else if (error.hasKeyNotInRegion()) {
+      // this error is reported from raftstore:
+      // key requested is not in current region
+      // should not happen here.
+      ByteString invalidKey = error.getKeyNotInRegion().getKey();
+      logger.error(
+          String.format(
+              "Key not in region [%s] for key [%s], this error should not happen here.",
+              recv.getRegion(), KeyUtils.formatBytesUTF8(invalidKey)));
+      regionManager.clearRegionCache();
+      throw new StatusRuntimeException(Status.UNKNOWN.withDescription(error.toString()));
+    }
+
+    logger.warn(String.format("Unknown error %s for region [%s]", error, recv.getRegion()));
+    // For other errors, we only drop cache here.
+    // Upper level may split this task.
+    invalidateRegionStoreCache(recv.getRegion());
+    // retry if raft proposal is dropped, it indicates the store is in the middle of transition
+    if (error.getMessage().contains("Raft ProposalDropped")) {
+      backOffer.doBackOff(
+          BackOffFunction.BackOffFuncType.BoRegionMiss, new GrpcException(error.getMessage()));
+      return true;
+    }
+    return false;
+  }
+
+  // ref: https://github.com/tikv/client-go/blob/tidb-5.2/internal/locate/region_request.go#L985
+  // OnRegionEpochNotMatch removes the old region and inserts new regions into the cache.
+  // It returns whether retries the request because it's possible the region epoch is ahead of
+  // TiKV's due to slow appling.
+  private boolean onRegionEpochNotMatch(BackOffer backOffer, List<Metapb.Region> currentRegions) {
+    if (currentRegions.size() == 0) {
+      this.regionManager.onRegionStale(recv.getRegion());
+      return false;
+    }
+
+    // Find whether the region epoch in `ctx` is ahead of TiKV's. If so, backoff.
+    for (Metapb.Region meta : currentRegions) {
+      if (meta.getId() == recv.getRegion().getId()
+          && (meta.getRegionEpoch().getConfVer() < recv.getRegion().getVerID().getConfVer()
+              || meta.getRegionEpoch().getVersion() < recv.getRegion().getVerID().getVer())) {
+        String errorMsg =
+            String.format(
+                "region epoch is ahead of tikv, region: %s, currentRegions: %s",
+                recv.getRegion(), currentRegions);
+        logger.info(errorMsg);
+        backOffer.doBackOff(
+            BackOffFunction.BackOffFuncType.BoRegionMiss, new TiKVException(errorMsg));
+        return true;
+      }
+    }
+
+    boolean needInvalidateOld = true;
+    List<TiRegion> newRegions = new ArrayList<>(currentRegions.size());
+    // If the region epoch is not ahead of TiKV's, replace region meta in region cache.
+    for (Metapb.Region meta : currentRegions) {
+      TiRegion region = regionManager.createRegion(meta, backOffer);
+      newRegions.add(region);
+      if (recv.getRegion().getVerID() == region.getVerID()) {
+        needInvalidateOld = false;
+      }
+    }
+
+    if (needInvalidateOld) {
+      this.regionManager.onRegionStale(recv.getRegion());
+    }
+
+    for (TiRegion region : newRegions) {
+      regionManager.insertRegionToCache(region);
+    }
+
+    return false;
+  }
+
+  @Override
+  public boolean handleRequestError(BackOffer backOffer, Exception e) {
+    if (recv.onStoreUnreachable(backOffer)) {
+      if (!backOffer.canRetryAfterSleep(BackOffFunction.BackOffFuncType.BoTiKVRPC)) {
+        regionManager.onRequestFail(recv.getRegion());
+        throw new GrpcException("retry is exhausted.", e);
+      }
+      return true;
+    }
+
+    logger.warn("request failed because of: " + e.getMessage());
+    if (!backOffer.canRetryAfterSleep(BackOffFunction.BackOffFuncType.BoTiKVRPC)) {
+      regionManager.onRequestFail(recv.getRegion());
+      throw new GrpcException(
+          "send tikv request error: " + e.getMessage() + ", try next peer later", e);
+    }
+    // TiKV maybe down, so do not retry in `callWithRetry`
+    // should re-fetch the new leader from PD and send request to it
+    return false;
+  }
+
+  public Errorpb.Error getRegionError(RespT resp) {
+    if (getRegionError != null) {
+      return getRegionError.apply(resp);
+    }
+    return null;
+  }
+
+  public TiRegion getRegion() {
+    return recv.getRegion();
+  }
+
+  private void invalidateRegionStoreCache(TiRegion ctxRegion) {
+    regionManager.invalidateRegion(ctxRegion);
+    regionManager.invalidateStore(ctxRegion.getLeader().getStoreId());
+  }
+}

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

@@ -25,10 +25,10 @@ import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
 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 RawScanIterator extends ScanIterator {
+  private final BackOffer scanBackOffer;
 
   public RawScanIterator(
       TiConfiguration conf,
@@ -36,15 +36,19 @@ public class RawScanIterator extends ScanIterator {
       ByteString startKey,
       ByteString endKey,
       int limit,
-      boolean keyOnly) {
+      boolean keyOnly,
+      BackOffer scanBackOffer) {
     super(conf, builder, startKey, endKey, limit, keyOnly);
+
+    this.scanBackOffer = scanBackOffer;
   }
 
+  @Override
   TiRegion loadCurrentRegionToCache() throws GrpcException {
-    BackOffer backOffer = ConcreteBackOffer.newScannerNextMaxBackOff();
+    BackOffer backOffer = scanBackOffer;
     while (true) {
-      try (RegionStoreClient client = builder.build(startKey)) {
-        client.setTimeout(conf.getScanTimeout());
+      try (RegionStoreClient client = builder.build(startKey, backOffer)) {
+        client.setTimeout(conf.getRawKVScanTimeoutInMS());
         TiRegion region = client.getRegion();
         if (limit <= 0) {
           currentCache = null;
@@ -78,26 +82,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;
-    // continue when cache is empty but not null
-    for (kv = getCurrent(); currentCache != null && kv == null; kv = getCurrent()) {
-      if (cacheLoadFails()) {
-        return null;
-      }
-    }
-    return kv;
+    return getCurrent();
   }
 }

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,9 +17,11 @@ package org.tikv.common.policy;
 
 import com.google.common.collect.ImmutableSet;
 import io.grpc.Status;
+import io.prometheus.client.Counter;
 import io.prometheus.client.Histogram;
 import java.util.concurrent.Callable;
 import org.tikv.common.exception.GrpcException;
+import org.tikv.common.log.SlowLogSpan;
 import org.tikv.common.operation.ErrorHandler;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ConcreteBackOffer;
@@ -32,6 +34,18 @@ public abstract class RetryPolicy<RespT> {
           .help("grpc request latency.")
           .labelNames("type")
           .register();
+  public static final Histogram CALL_WITH_RETRY_DURATION =
+      Histogram.build()
+          .name("client_java_call_with_retry_duration")
+          .help("callWithRetry duration.")
+          .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;
@@ -54,35 +68,50 @@ public abstract class RetryPolicy<RespT> {
     }
   }
 
-  public RespT callWithRetry(Callable<RespT> proc, String methodName) {
-    while (true) {
-      RespT result = null;
-      try {
-        // add single request duration histogram
-        Histogram.Timer requestTimer = GRPC_SINGLE_REQUEST_LATENCY.labels(methodName).startTimer();
+  public RespT callWithRetry(Callable<RespT> proc, String methodName, BackOffer backOffer) {
+    Histogram.Timer callWithRetryTimer = CALL_WITH_RETRY_DURATION.labels(methodName).startTimer();
+    SlowLogSpan callWithRetrySlowLogSpan =
+        backOffer.getSlowLog().start("callWithRetry " + methodName);
+    try {
+      while (true) {
+        RespT result = null;
         try {
-          result = proc.call();
-        } finally {
-          requestTimer.observeDuration();
+          // add single request duration histogram
+          Histogram.Timer requestTimer =
+              GRPC_SINGLE_REQUEST_LATENCY.labels(methodName).startTimer();
+          SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("gRPC " + methodName);
+          try {
+            result = proc.call();
+          } finally {
+            slowLogSpan.end();
+            requestTimer.observeDuration();
+          }
+        } catch (Exception e) {
+          rethrowNotRecoverableException(e);
+          // Handle request call error
+          backOffer.checkTimeout();
+          boolean retry = handler.handleRequestError(backOffer, e);
+          if (retry) {
+            GRPC_REQUEST_RETRY_NUM.labels(methodName).inc();
+            continue;
+          } else {
+            return result;
+          }
         }
-      } catch (Exception e) {
-        rethrowNotRecoverableException(e);
-        // Handle request call error
-        boolean retry = handler.handleRequestError(backOffer, e);
-        if (retry) {
-          continue;
-        }
-      }
 
-      // Handle response error
-      if (handler != null) {
-        boolean retry = handler.handleResponseError(backOffer, result);
-        if (retry) {
-          // add retry counter
-          continue;
+        // Handle response error
+        if (handler != null) {
+          boolean retry = handler.handleResponseError(backOffer, result);
+          if (retry) {
+            GRPC_REQUEST_RETRY_NUM.labels(methodName).inc();
+            continue;
+          }
         }
+        return result;
       }
-      return result;
+    } finally {
+      callWithRetryTimer.observeDuration();
+      callWithRetrySlowLogSpan.end();
     }
   }
 

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

@@ -20,28 +20,56 @@ package org.tikv.common.region;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.protobuf.ByteString;
 import io.grpc.ManagedChannel;
-import java.util.concurrent.TimeUnit;
+import io.grpc.Metadata;
+import io.grpc.stub.MetadataUtils;
+import io.prometheus.client.Histogram;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.*;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.tikv.common.AbstractGRPCClient;
 import org.tikv.common.TiConfiguration;
 import org.tikv.common.exception.GrpcException;
+import org.tikv.common.log.SlowLogSpan;
+import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
+import org.tikv.kvproto.Kvrpcpb;
 import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.TikvGrpc;
 
 public abstract class AbstractRegionStoreClient
-    extends AbstractGRPCClient<TikvGrpc.TikvBlockingStub, TikvGrpc.TikvStub>
+    extends AbstractGRPCClient<TikvGrpc.TikvBlockingStub, TikvGrpc.TikvFutureStub>
     implements RegionErrorReceiver {
+  private static final Logger logger = LoggerFactory.getLogger(AbstractRegionStoreClient.class);
+
+  public static final Histogram SEEK_LEADER_STORE_DURATION =
+      Histogram.build()
+          .name("client_java_seek_leader_store_duration")
+          .help("seek leader store duration.")
+          .register();
+
+  public static final Histogram SEEK_PROXY_STORE_DURATION =
+      Histogram.build()
+          .name("client_java_seek_proxy_store_duration")
+          .help("seek proxy store duration.")
+          .register();
 
   protected final RegionManager regionManager;
   protected TiRegion region;
+  protected TiStore store;
 
   protected AbstractRegionStoreClient(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       ChannelFactory channelFactory,
       TikvGrpc.TikvBlockingStub blockingStub,
-      TikvGrpc.TikvStub asyncStub,
+      TikvGrpc.TikvFutureStub asyncStub,
       RegionManager regionManager) {
     super(conf, channelFactory, blockingStub, asyncStub);
     checkNotNull(region, "Region is empty");
@@ -49,8 +77,13 @@ public abstract class AbstractRegionStoreClient
     checkArgument(region.getLeader() != null, "Leader Peer is null");
     this.region = region;
     this.regionManager = regionManager;
+    this.store = store;
+    if (this.store.getProxyStore() != null) {
+      this.timeout = conf.getForwardTimeout();
+    }
   }
 
+  @Override
   public TiRegion getRegion() {
     return region;
   }
@@ -61,7 +94,7 @@ public abstract class AbstractRegionStoreClient
   }
 
   @Override
-  protected TikvGrpc.TikvStub getAsyncStub() {
+  protected TikvGrpc.TikvFutureStub getAsyncStub() {
     return asyncStub.withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
   }
 
@@ -71,43 +104,271 @@ 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.
     if (!region.getRegionEpoch().equals(newRegion.getRegionEpoch())) {
       return false;
     }
+
+    // If we try one peer but find the leader has not changed, we do not need to try other peers.
+    if (region.getLeader().getStoreId() == newRegion.getLeader().getStoreId()) {
+      store = null;
+    }
     region = newRegion;
-    String addressStr = regionManager.getStoreById(region.getLeader().getStoreId()).getAddress();
-    ManagedChannel channel =
-        channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
-    blockingStub = TikvGrpc.newBlockingStub(channel);
-    asyncStub = TikvGrpc.newStub(channel);
+    store = regionManager.getStoreById(region.getLeader().getStoreId());
+    updateClientStub();
     return true;
   }
 
   @Override
-  public void onStoreNotMatch(Metapb.Store store) {
-    String addressStr = store.getAddress();
+  public boolean onStoreUnreachable(BackOffer backOffer) {
+    if (!store.isValid()) {
+      logger.warn(String.format("store [%d] has been invalid", store.getId()));
+      store = regionManager.getStoreById(store.getId(), backOffer);
+      updateClientStub();
+      return true;
+    }
+
+    // seek an available leader store to send request
+    backOffer.checkTimeout();
+    Boolean result = seekLeaderStore(backOffer);
+    if (result != null) {
+      return result;
+    }
+    if (conf.getEnableGrpcForward()) {
+      // seek an available proxy store to forward request
+      backOffer.checkTimeout();
+      return seekProxyStore(backOffer);
+    }
+    return false;
+  }
+
+  protected Kvrpcpb.Context makeContext(TiStoreType storeType) {
+    return region.getReplicaContext(java.util.Collections.emptySet(), storeType);
+  }
+
+  protected Kvrpcpb.Context makeContext(Set<Long> resolvedLocks, TiStoreType storeType) {
+    return region.getReplicaContext(resolvedLocks, storeType);
+  }
+
+  private void updateClientStub() {
+    String addressStr = store.getStore().getAddress();
+    long deadline = timeout;
+    if (store.getProxyStore() != null) {
+      addressStr = store.getProxyStore().getAddress();
+      deadline = conf.getForwardTimeout();
+    }
     ManagedChannel channel =
         channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
-    blockingStub = TikvGrpc.newBlockingStub(channel);
-    asyncStub = TikvGrpc.newStub(channel);
-    if (region.getLeader().getStoreId() != store.getId()) {
-      logger.warn(
-          "store_not_match may occur? "
-              + region
-              + ", original store = "
-              + store.getId()
-              + " address = "
-              + addressStr);
+    blockingStub =
+        TikvGrpc.newBlockingStub(channel).withDeadlineAfter(deadline, TimeUnit.MILLISECONDS);
+    asyncStub = TikvGrpc.newFutureStub(channel).withDeadlineAfter(deadline, TimeUnit.MILLISECONDS);
+    if (store.getProxyStore() != null) {
+      Metadata header = new Metadata();
+      header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
+      blockingStub = MetadataUtils.attachHeaders(blockingStub, header);
+      asyncStub = MetadataUtils.attachHeaders(asyncStub, header);
+    }
+  }
+
+  private Boolean seekLeaderStore(BackOffer backOffer) {
+    Histogram.Timer switchLeaderDurationTimer = SEEK_LEADER_STORE_DURATION.startTimer();
+    SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("seekLeaderStore");
+    try {
+      List<Metapb.Peer> peers = region.getFollowerList();
+      if (peers.isEmpty()) {
+        // no followers available, retry
+        logger.warn(String.format("no followers of region[%d] available, retry", region.getId()));
+        regionManager.onRequestFail(region);
+        return false;
+      }
+
+      logger.info(String.format("try switch leader: region[%d]", region.getId()));
+
+      Metapb.Peer peer = switchLeaderStore();
+      if (peer != null) {
+        // we found a leader
+        TiStore currentLeaderStore = regionManager.getStoreById(peer.getStoreId());
+        if (currentLeaderStore.isReachable()) {
+          logger.info(
+              String.format(
+                  "update leader using switchLeader logic from store[%d] to store[%d]",
+                  region.getLeader().getStoreId(), peer.getStoreId()));
+          // update region cache
+          TiRegion result = regionManager.updateLeader(region, peer.getStoreId());
+          if (result != null) {
+            region = result;
+            // switch to leader store
+            store = currentLeaderStore;
+            updateClientStub();
+          }
+          return false;
+        }
+      } else {
+        // no leader found, some response does not return normally, there may be network partition.
+        logger.warn(
+            String.format(
+                "leader for region[%d] is not found, it is possible that network partition occurred",
+                region.getId()));
+      }
+    } finally {
+      switchLeaderDurationTimer.observeDuration();
+      slowLogSpan.end();
+    }
+    return null;
+  }
+
+  private boolean seekProxyStore(BackOffer backOffer) {
+    SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("seekProxyStore");
+    Histogram.Timer grpcForwardDurationTimer = SEEK_PROXY_STORE_DURATION.startTimer();
+    try {
+      logger.info(String.format("try grpc forward: region[%d]", region.getId()));
+      // when current leader cannot be reached
+      TiStore storeWithProxy = switchProxyStore();
+      if (storeWithProxy == null) {
+        // no store available, retry
+        logger.warn(String.format("No store available, retry: region[%d]", region.getId()));
+        return false;
+      }
+      // use proxy store to forward requests
+      regionManager.updateStore(store, storeWithProxy);
+      store = storeWithProxy;
+      updateClientStub();
+      return true;
+    } finally {
+      grpcForwardDurationTimer.observeDuration();
+      slowLogSpan.end();
+    }
+  }
+
+  // first: leader peer, second: true if any responses returned with grpc error
+  private Metapb.Peer switchLeaderStore() {
+    List<SwitchLeaderTask> responses = new LinkedList<>();
+    for (Metapb.Peer peer : region.getFollowerList()) {
+      ByteString key = region.getStartKey();
+      TiStore peerStore = regionManager.getStoreById(peer.getStoreId());
+      ManagedChannel channel =
+          channelFactory.getChannel(
+              peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
+      TikvGrpc.TikvFutureStub stub =
+          TikvGrpc.newFutureStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
+      Kvrpcpb.RawGetRequest rawGetRequest =
+          Kvrpcpb.RawGetRequest.newBuilder()
+              .setContext(region.getReplicaContext(peer))
+              .setKey(key)
+              .build();
+      ListenableFuture<Kvrpcpb.RawGetResponse> task = stub.rawGet(rawGetRequest);
+      responses.add(new SwitchLeaderTask(task, peer));
+    }
+    while (true) {
+      try {
+        Thread.sleep(2);
+      } catch (InterruptedException e) {
+        throw new GrpcException(e);
+      }
+      List<SwitchLeaderTask> unfinished = new LinkedList<>();
+      for (SwitchLeaderTask task : responses) {
+        if (!task.task.isDone()) {
+          unfinished.add(task);
+          continue;
+        }
+        try {
+          Kvrpcpb.RawGetResponse resp = task.task.get();
+          if (resp != null) {
+            if (!resp.hasRegionError()) {
+              // the peer is leader
+              logger.info(
+                  String.format("rawGet response indicates peer[%d] is leader", task.peer.getId()));
+              return task.peer;
+            }
+          }
+        } catch (Exception ignored) {
+        }
+      }
+      if (unfinished.isEmpty()) {
+        return null;
+      }
+      responses = unfinished;
+    }
+  }
+
+  private TiStore switchProxyStore() {
+    long forwardTimeout = conf.getForwardTimeout();
+    List<ForwardCheckTask> responses = new LinkedList<>();
+    for (Metapb.Peer peer : region.getFollowerList()) {
+      ByteString key = region.getStartKey();
+      TiStore peerStore = regionManager.getStoreById(peer.getStoreId());
+      ManagedChannel channel =
+          channelFactory.getChannel(
+              peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
+      TikvGrpc.TikvFutureStub stub =
+          TikvGrpc.newFutureStub(channel).withDeadlineAfter(forwardTimeout, TimeUnit.MILLISECONDS);
+      Metadata header = new Metadata();
+      header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
+      Kvrpcpb.RawGetRequest rawGetRequest =
+          Kvrpcpb.RawGetRequest.newBuilder()
+              .setContext(region.getReplicaContext(peer))
+              .setKey(key)
+              .build();
+      ListenableFuture<Kvrpcpb.RawGetResponse> task =
+          MetadataUtils.attachHeaders(stub, header).rawGet(rawGetRequest);
+      responses.add(new ForwardCheckTask(task, peerStore.getStore()));
+    }
+    while (true) {
+      try {
+        Thread.sleep(2);
+      } catch (InterruptedException e) {
+        throw new GrpcException(e);
+      }
+      List<ForwardCheckTask> unfinished = new LinkedList<>();
+      for (ForwardCheckTask task : responses) {
+        if (!task.task.isDone()) {
+          unfinished.add(task);
+          continue;
+        }
+        try {
+          // any answer will do
+          Kvrpcpb.RawGetResponse resp = task.task.get();
+          logger.info(
+              String.format(
+                  "rawGetResponse indicates forward from [%s] to [%s]",
+                  task.store.getAddress(), store.getAddress()));
+          return store.withProxy(task.store);
+        } catch (Exception ignored) {
+        }
+      }
+      if (unfinished.isEmpty()) {
+        return null;
+      }
+      responses = unfinished;
+    }
+  }
+
+  private static class SwitchLeaderTask {
+    private final ListenableFuture<Kvrpcpb.RawGetResponse> task;
+    private final Metapb.Peer peer;
+
+    private SwitchLeaderTask(ListenableFuture<Kvrpcpb.RawGetResponse> task, Metapb.Peer peer) {
+      this.task = task;
+      this.peer = peer;
+    }
+  }
+
+  private static class ForwardCheckTask {
+    private final ListenableFuture<Kvrpcpb.RawGetResponse> task;
+    private final Metapb.Store store;
+
+    private ForwardCheckTask(ListenableFuture<Kvrpcpb.RawGetResponse> task, Metapb.Store store) {
+      this.task = task;
+      this.store = store;
     }
   }
 }

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

@@ -0,0 +1,173 @@
+package org.tikv.common.region;
+
+import static org.tikv.common.codec.KeyUtils.formatBytesUTF8;
+import static org.tikv.common.util.KeyRangeUtils.makeRange;
+
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.TreeRangeMap;
+import com.google.protobuf.ByteString;
+import java.util.HashMap;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.key.Key;
+import org.tikv.common.util.BackOffer;
+
+public class RegionCache {
+  private static final Logger logger = LoggerFactory.getLogger(RegionCache.class);
+
+  private final Map<Long, TiRegion> regionCache;
+  private final Map<Long, TiStore> storeCache;
+  private final RangeMap<Key, Long> keyToRegionIdCache;
+
+  public RegionCache() {
+    regionCache = new HashMap<>();
+    storeCache = new HashMap<>();
+
+    keyToRegionIdCache = TreeRangeMap.create();
+  }
+
+  public synchronized TiRegion getRegionByKey(ByteString key, BackOffer backOffer) {
+    Long regionId;
+    if (key.isEmpty()) {
+      // if key is empty, it must be the start key.
+      regionId = keyToRegionIdCache.get(Key.toRawKey(key, true));
+    } else {
+      regionId = keyToRegionIdCache.get(Key.toRawKey(key));
+    }
+    if (logger.isDebugEnabled()) {
+      logger.debug(
+          String.format("getRegionByKey key[%s] -> ID[%s]", formatBytesUTF8(key), regionId));
+    }
+
+    if (regionId == null) {
+      return null;
+    }
+    TiRegion region;
+    region = regionCache.get(regionId);
+    if (logger.isDebugEnabled()) {
+      logger.debug(String.format("getRegionByKey ID[%s] -> Region[%s]", regionId, region));
+    }
+    return region;
+  }
+
+  public synchronized TiRegion putRegion(TiRegion region) {
+    if (logger.isDebugEnabled()) {
+      logger.debug("putRegion: " + region);
+    }
+    TiRegion oldRegion = regionCache.get(region.getId());
+    if (oldRegion != null) {
+      if (oldRegion.getMeta().equals(region.getMeta())) {
+        return oldRegion;
+      } else {
+        invalidateRegion(oldRegion);
+      }
+    }
+    regionCache.put(region.getId(), region);
+    keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
+    return region;
+  }
+
+  @Deprecated
+  public synchronized TiRegion getRegionById(long regionId) {
+    TiRegion region = regionCache.get(regionId);
+    if (logger.isDebugEnabled()) {
+      logger.debug(String.format("getRegionByKey ID[%s] -> Region[%s]", regionId, region));
+    }
+    return region;
+  }
+
+  /** Removes region associated with regionId from regionCache. */
+  public synchronized void invalidateRegion(TiRegion region) {
+    try {
+      if (logger.isDebugEnabled()) {
+        logger.debug(String.format("invalidateRegion ID[%s]", region.getId()));
+      }
+      TiRegion oldRegion = regionCache.get(region.getId());
+      if (oldRegion != null && oldRegion == region) {
+        keyToRegionIdCache.remove(makeRange(region.getStartKey(), region.getEndKey()));
+        regionCache.remove(region.getId());
+      }
+    } catch (Exception ignore) {
+    }
+  }
+
+  public synchronized void insertRegionToCache(TiRegion region) {
+    try {
+      TiRegion oldRegion = regionCache.get(region.getId());
+      if (oldRegion != null) {
+        keyToRegionIdCache.remove(makeRange(oldRegion.getStartKey(), oldRegion.getEndKey()));
+      }
+      regionCache.put(region.getId(), region);
+      keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
+    } catch (Exception ignore) {
+    }
+  }
+
+  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.getMeta().equals(oldRegion.getMeta())) {
+        return false;
+      } else {
+        if (oldRegion != null) {
+          keyToRegionIdCache.remove(makeRange(oldRegion.getStartKey(), oldRegion.getEndKey()));
+        }
+        regionCache.put(region.getId(), region);
+        keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
+        return true;
+      }
+    } catch (Exception ignore) {
+      return false;
+    }
+  }
+
+  public synchronized boolean updateStore(TiStore oldStore, TiStore newStore) {
+    if (!newStore.isValid()) {
+      return false;
+    }
+    if (oldStore == null) {
+      storeCache.put(newStore.getId(), newStore);
+      return true;
+    }
+    TiStore originStore = storeCache.get(oldStore.getId());
+    if (originStore.equals(oldStore)) {
+      storeCache.put(newStore.getId(), newStore);
+      oldStore.markInvalid();
+      return true;
+    }
+    return false;
+  }
+
+  public synchronized void invalidateStore(long storeId) {
+    TiStore store = storeCache.remove(storeId);
+    if (store != null) {
+      store.markInvalid();
+    }
+  }
+
+  public synchronized TiStore getStoreById(long id) {
+    return storeCache.get(id);
+  }
+
+  public synchronized boolean putStore(long id, TiStore store) {
+    TiStore oldStore = storeCache.get(id);
+    if (oldStore != null) {
+      if (oldStore.equals(store)) {
+        return false;
+      } else {
+        oldStore.markInvalid();
+      }
+    }
+    storeCache.put(id, store);
+    return true;
+  }
+
+  public synchronized void clearAll() {
+    keyToRegionIdCache.clear();
+    regionCache.clear();
+  }
+}

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

@@ -17,12 +17,13 @@
 
 package org.tikv.common.region;
 
-import org.tikv.kvproto.Metapb.Store;
+import org.tikv.common.util.BackOffer;
 
 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(BackOffer backOffer);
 
   TiRegion getRegion();
 }

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

@@ -18,77 +18,100 @@
 package org.tikv.common.region;
 
 import static org.tikv.common.codec.KeyUtils.formatBytesUTF8;
-import static org.tikv.common.util.KeyRangeUtils.makeRange;
 
-import com.google.common.collect.RangeMap;
-import com.google.common.collect.TreeRangeMap;
 import com.google.protobuf.ByteString;
 import io.prometheus.client.Histogram;
-import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.function.Function;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.ReadOnlyPDClient;
-import org.tikv.common.event.CacheInvalidateEvent;
+import org.tikv.common.TiConfiguration;
 import org.tikv.common.exception.GrpcException;
+import org.tikv.common.exception.InvalidStoreException;
 import org.tikv.common.exception.TiClientInternalException;
-import org.tikv.common.key.Key;
+import org.tikv.common.log.SlowLogSpan;
 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")
 public class RegionManager {
   private static final Logger logger = LoggerFactory.getLogger(RegionManager.class);
-  // TODO: the region cache logic need rewrite.
-  // https://github.com/pingcap/tispark/issues/1170
-  private final RegionCache cache;
-  private final boolean isReplicaRead;
-
-  private final Function<CacheInvalidateEvent, Void> cacheInvalidateCallback;
-
   public static final Histogram GET_REGION_BY_KEY_REQUEST_LATENCY =
       Histogram.build()
           .name("client_java_get_region_by_requests_latency")
           .help("getRegionByKey request latency.")
           .register();
 
-  // To avoid double retrieval, we used the async version of grpc
-  // When rpc not returned, instead of call again, it wait for previous one done
+  // TODO: the region cache logic need rewrite.
+  // https://github.com/pingcap/tispark/issues/1170
+  private final RegionCache cache;
+  private final ReadOnlyPDClient pdClient;
+  private final TiConfiguration conf;
+  private final ScheduledExecutorService executor;
+  private final StoreHealthyChecker storeChecker;
+
   public RegionManager(
-      ReadOnlyPDClient pdClient, Function<CacheInvalidateEvent, Void> cacheInvalidateCallback) {
-    this.cache = new RegionCache(pdClient);
-    this.isReplicaRead = pdClient.isReplicaRead();
-    this.cacheInvalidateCallback = cacheInvalidateCallback;
+      TiConfiguration conf, ReadOnlyPDClient pdClient, ChannelFactory channelFactory) {
+    this.cache = new RegionCache();
+    this.pdClient = pdClient;
+    this.conf = conf;
+    long period = conf.getHealthCheckPeriodDuration();
+    StoreHealthyChecker storeChecker =
+        new StoreHealthyChecker(
+            channelFactory, pdClient, this.cache, conf.getGrpcHealthCheckTimeout());
+    this.storeChecker = storeChecker;
+    this.executor = Executors.newScheduledThreadPool(1);
+    this.executor.scheduleAtFixedRate(storeChecker, period, period, TimeUnit.MILLISECONDS);
   }
 
-  public RegionManager(ReadOnlyPDClient pdClient) {
-    this.cache = new RegionCache(pdClient);
-    this.isReplicaRead = pdClient.isReplicaRead();
-    this.cacheInvalidateCallback = null;
+  public RegionManager(TiConfiguration conf, ReadOnlyPDClient pdClient) {
+    this.cache = new RegionCache();
+    this.pdClient = pdClient;
+    this.conf = conf;
+    this.storeChecker = null;
+    this.executor = null;
   }
 
-  public Function<CacheInvalidateEvent, Void> getCacheInvalidateCallback() {
-    return cacheInvalidateCallback;
+  public synchronized void close() {
+    if (this.executor != null) {
+      this.executor.shutdownNow();
+    }
   }
 
   public ReadOnlyPDClient getPDClient() {
-    return this.cache.pdClient;
+    return this.pdClient;
   }
 
   public TiRegion getRegionByKey(ByteString key) {
-    return getRegionByKey(key, ConcreteBackOffer.newGetBackOff());
+    return getRegionByKey(key, defaultBackOff());
   }
 
   public TiRegion getRegionByKey(ByteString key, BackOffer backOffer) {
-    return cache.getRegionByKey(key, backOffer);
+    Histogram.Timer requestTimer = GET_REGION_BY_KEY_REQUEST_LATENCY.startTimer();
+    SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("getRegionByKey");
+    TiRegion region = cache.getRegionByKey(key, backOffer);
+    try {
+      if (region == null) {
+        logger.debug("Key not found in keyToRegionIdCache:" + formatBytesUTF8(key));
+        Pair<Metapb.Region, Metapb.Peer> regionAndLeader = pdClient.getRegionByKey(backOffer, key);
+        region =
+            cache.putRegion(createRegion(regionAndLeader.first, regionAndLeader.second, backOffer));
+      }
+    } finally {
+      requestTimer.observeDuration();
+      slowLogSpan.end();
+    }
+
+    return region;
   }
 
   @Deprecated
@@ -99,45 +122,45 @@ public class RegionManager {
   // Consider region A, B. After merge of (A, B) -> A, region ID B does not exist.
   // This request is unrecoverable.
   public TiRegion getRegionById(long regionId) {
-    return cache.getRegionById(ConcreteBackOffer.newGetBackOff(), regionId);
+    BackOffer backOffer = defaultBackOff();
+    TiRegion region = cache.getRegionById(regionId);
+    if (region == null) {
+      Pair<Metapb.Region, Metapb.Peer> regionAndLeader =
+          pdClient.getRegionByID(backOffer, regionId);
+      region = createRegion(regionAndLeader.first, regionAndLeader.second, backOffer);
+      return cache.putRegion(region);
+    }
+    return region;
   }
 
-  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) {
-    return getRegionStorePairByKey(key, storeType, ConcreteBackOffer.newGetBackOff());
+  public Pair<TiRegion, TiStore> getRegionStorePairByKey(ByteString key, TiStoreType storeType) {
+    return getRegionStorePairByKey(key, storeType, defaultBackOff());
   }
 
-  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) {
-      throw new TiClientInternalException("Region not exist for key:" + formatBytesUTF8(key));
-    }
+    TiRegion region = getRegionByKey(key, backOffer);
     if (!region.isValid()) {
       throw new TiClientInternalException("Region invalid: " + region.toString());
     }
 
-    Store store = null;
+    TiStore store = null;
     if (storeType == TiStoreType.TiKV) {
-      if (isReplicaRead) {
-        Peer peer = region.getCurrentFollower();
-        store = cache.getStoreById(peer.getStoreId(), backOffer);
-      } else {
-        Peer leader = region.getLeader();
-        store = cache.getStoreById(leader.getStoreId(), backOffer);
-      }
+      Peer peer = region.getCurrentReplica();
+      store = getStoreById(peer.getStoreId(), backOffer);
     } else {
       outerLoop:
       for (Peer peer : region.getLearnerList()) {
-        Store s = getStoreById(peer.getStoreId(), backOffer);
-        for (Metapb.StoreLabel label : s.getLabelsList()) {
+        TiStore s = getStoreById(peer.getStoreId(), backOffer);
+        for (Metapb.StoreLabel label : s.getStore().getLabelsList()) {
           if (label.getKey().equals(storeType.getLabelKey())
               && label.getValue().equals(storeType.getLabelValue())) {
             store = s;
@@ -146,61 +169,110 @@ public class RegionManager {
         }
       }
       if (store == null) {
-        // clear the region cache so we may get the learner peer next time
+        // clear the region cache, so we may get the learner peer next time
         cache.invalidateRegion(region);
       }
     }
 
-    if (store == null) {
-      throw new TiClientInternalException(
-          "Cannot find valid store on " + storeType + " for region " + region.toString());
-    }
-
     return Pair.create(region, store);
   }
 
-  public Store getStoreById(long id) {
-    return getStoreById(id, ConcreteBackOffer.newGetBackOff());
+  public TiRegion createRegion(Metapb.Region region, BackOffer backOffer) {
+    List<Metapb.Peer> peers = region.getPeersList();
+    List<TiStore> stores = getRegionStore(peers, backOffer);
+    return new TiRegion(conf, region, null, peers, stores);
   }
 
-  public Store getStoreById(long id, BackOffer backOffer) {
-    return cache.getStoreById(id, backOffer);
+  private TiRegion createRegion(Metapb.Region region, Metapb.Peer leader, BackOffer backOffer) {
+    List<Metapb.Peer> peers = region.getPeersList();
+    List<TiStore> stores = getRegionStore(peers, backOffer);
+    return new TiRegion(conf, region, leader, peers, stores);
+  }
+
+  private List<TiStore> getRegionStore(List<Metapb.Peer> peers, BackOffer backOffer) {
+    return peers
+        .stream()
+        .map(p -> getStoreById(p.getStoreId(), backOffer))
+        .collect(Collectors.toList());
+  }
+
+  private TiStore getStoreByIdWithBackOff(long id, BackOffer backOffer) {
+    try {
+      TiStore store = cache.getStoreById(id);
+      if (store == null) {
+        store = new TiStore(pdClient.getStore(backOffer, id));
+      } else {
+        return store;
+      }
+      // if we did not get store info from pd, remove store from cache
+      if (store.getStore() == null) {
+        logger.warn(String.format("failed to get store %d from pd", id));
+        return null;
+      }
+      // if the store is already tombstone, remove store from cache
+      if (store.getStore().getState().equals(StoreState.Tombstone)) {
+        logger.warn(String.format("store %d is tombstone", id));
+        return null;
+      }
+      if (cache.putStore(id, store) && storeChecker != null) {
+        storeChecker.scheduleStoreHealthCheck(store);
+      }
+      return store;
+    } catch (Exception e) {
+      throw new GrpcException(e);
+    }
+  }
+
+  public TiStore getStoreById(long id) {
+    return getStoreById(id, defaultBackOff());
+  }
+
+  public TiStore getStoreById(long id, BackOffer backOffer) {
+    TiStore store = getStoreByIdWithBackOff(id, backOffer);
+    if (store == null) {
+      logger.warn(String.format("failed to fetch store %d, the store may be missing", id));
+      cache.clearAll();
+      throw new InvalidStoreException(id);
+    }
+    return store;
   }
 
   public void onRegionStale(TiRegion region) {
     cache.invalidateRegion(region);
   }
 
-  public synchronized TiRegion updateLeader(TiRegion region, long storeId) {
-    TiRegion r = cache.getRegionFromCache(region.getId());
-    if (r != null) {
-      if (r.getLeader().getStoreId() == storeId) {
-        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 + ")");
+  public TiRegion updateLeader(TiRegion region, long storeId) {
+    if (region.getLeader().getStoreId() == storeId) {
+      return region;
     }
+    TiRegion newRegion = region.switchPeer(storeId);
+    if (cache.updateRegion(region, 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 + ")");
     return null;
   }
 
+  public synchronized void updateStore(TiStore oldStore, TiStore newStore) {
+    if (cache.updateStore(oldStore, newStore) && storeChecker != null) {
+      storeChecker.scheduleStoreHealthCheck(newStore);
+    }
+  }
+
+  /** Clears all cache when some unexpected error occurs. */
+  public void clearRegionCache() {
+    cache.clearAll();
+  }
+
   /**
    * Clears all cache when a TiKV server does not respond
    *
    * @param region region
    */
-  public void onRequestFail(TiRegion region) {
-    onRequestFail(region, region.getLeader().getStoreId());
-  }
-
-  private void onRequestFail(TiRegion region, long storeId) {
+  public synchronized void onRequestFail(TiRegion region) {
     cache.invalidateRegion(region);
-    cache.invalidateAllRegionForStore(storeId);
   }
 
   public void invalidateStore(long storeId) {
@@ -211,134 +283,11 @@ public class RegionManager {
     cache.invalidateRegion(region);
   }
 
-  public static class RegionCache {
-    private final Map<Long, TiRegion> regionCache;
-    private final Map<Long, Store> storeCache;
-    private final RangeMap<Key, Long> keyToRegionIdCache;
-    private final ReadOnlyPDClient pdClient;
+  public void insertRegionToCache(TiRegion region) {
+    cache.insertRegionToCache(region);
+  }
 
-    public RegionCache(ReadOnlyPDClient pdClient) {
-      regionCache = new HashMap<>();
-      storeCache = new HashMap<>();
-
-      keyToRegionIdCache = TreeRangeMap.create();
-      this.pdClient = pdClient;
-    }
-
-    public synchronized TiRegion getRegionByKey(ByteString key, BackOffer backOffer) {
-      Histogram.Timer requestTimer = GET_REGION_BY_KEY_REQUEST_LATENCY.startTimer();
-      try {
-        Long regionId;
-        if (key.isEmpty()) {
-          // if key is empty, it must be the start key.
-          regionId = keyToRegionIdCache.get(Key.toRawKey(key, true));
-        } else {
-          regionId = keyToRegionIdCache.get(Key.toRawKey(key));
-        }
-        if (logger.isDebugEnabled()) {
-          logger.debug(
-              String.format("getRegionByKey key[%s] -> ID[%s]", formatBytesUTF8(key), regionId));
-        }
-
-        if (regionId == null) {
-          logger.debug("Key not found in keyToRegionIdCache:" + formatBytesUTF8(key));
-          TiRegion region = pdClient.getRegionByKey(backOffer, key);
-          if (!putRegion(region)) {
-            throw new TiClientInternalException("Invalid Region: " + region.toString());
-          }
-          return region;
-        }
-        TiRegion region;
-        region = regionCache.get(regionId);
-        if (logger.isDebugEnabled()) {
-          logger.debug(String.format("getRegionByKey ID[%s] -> Region[%s]", regionId, region));
-        }
-
-        return region;
-      } finally {
-        requestTimer.observeDuration();
-      }
-    }
-
-    private synchronized boolean putRegion(TiRegion region) {
-      if (logger.isDebugEnabled()) {
-        logger.debug("putRegion: " + region);
-      }
-      regionCache.put(region.getId(), region);
-      keyToRegionIdCache.put(makeRange(region.getStartKey(), region.getEndKey()), region.getId());
-      return true;
-    }
-
-    @Deprecated
-    private synchronized TiRegion getRegionById(BackOffer backOffer, long regionId) {
-      TiRegion region = regionCache.get(regionId);
-      if (logger.isDebugEnabled()) {
-        logger.debug(String.format("getRegionByKey ID[%s] -> Region[%s]", regionId, region));
-      }
-      if (region == null) {
-        region = pdClient.getRegionByID(backOffer, regionId);
-        if (!putRegion(region)) {
-          throw new TiClientInternalException("Invalid Region: " + region.toString());
-        }
-      }
-      return region;
-    }
-
-    private synchronized TiRegion getRegionFromCache(long regionId) {
-      return regionCache.get(regionId);
-    }
-
-    /** Removes region associated with regionId from regionCache. */
-    public synchronized void invalidateRegion(TiRegion region) {
-      try {
-        if (logger.isDebugEnabled()) {
-          logger.debug(String.format("invalidateRegion ID[%s]", region.getId()));
-        }
-        TiRegion oldRegion = regionCache.get(region.getId());
-        if (oldRegion != null && oldRegion == region) {
-          keyToRegionIdCache.remove(makeRange(region.getStartKey(), region.getEndKey()));
-          regionCache.remove(region.getId());
-        }
-      } catch (Exception ignore) {
-      }
-    }
-
-    public synchronized void invalidateAllRegionForStore(long storeId) {
-      List<TiRegion> regionToRemove = new ArrayList<>();
-      for (TiRegion r : regionCache.values()) {
-        if (r.getLeader().getStoreId() == storeId) {
-          if (logger.isDebugEnabled()) {
-            logger.debug(String.format("invalidateAllRegionForStore Region[%s]", r));
-          }
-          regionToRemove.add(r);
-        }
-      }
-
-      // remove region
-      for (TiRegion r : regionToRemove) {
-        regionCache.remove(r.getId());
-        keyToRegionIdCache.remove(makeRange(r.getStartKey(), r.getEndKey()));
-      }
-    }
-
-    public synchronized void invalidateStore(long storeId) {
-      storeCache.remove(storeId);
-    }
-
-    public synchronized Store getStoreById(long id, BackOffer backOffer) {
-      try {
-        Store store = storeCache.get(id);
-        if (store == null) {
-          store = pdClient.getStore(backOffer, id);
-        }
-        if (store.getState().equals(StoreState.Tombstone)) {
-          return null;
-        }
-        storeCache.put(id, store);
-        return store;
-      } catch (Exception e) {
-        throw new GrpcException(e);
-      }
-    }
+  private BackOffer defaultBackOff() {
+    return ConcreteBackOffer.newCustomBackOff(conf.getRawKVDefaultBackoffInMS());
   }
 }

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

@@ -26,10 +26,11 @@ 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;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.PDClient;
@@ -38,15 +39,16 @@ import org.tikv.common.TiConfiguration;
 import org.tikv.common.Version;
 import org.tikv.common.exception.*;
 import org.tikv.common.operation.KVErrorHandler;
+import org.tikv.common.operation.RegionErrorHandler;
 import org.tikv.common.streaming.StreamingResponse;
 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.Metapb;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
-import org.tikv.kvproto.TikvGrpc.TikvStub;
+import org.tikv.kvproto.TikvGrpc.TikvFutureStub;
 import org.tikv.txn.AbstractLockResolverClient;
 import org.tikv.txn.Lock;
 import org.tikv.txn.ResolveLockResult;
@@ -87,23 +89,23 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
   private RegionStoreClient(
       TiConfiguration conf,
       TiRegion region,
-      String storeVersion,
+      TiStore store,
       TiStoreType storeType,
       ChannelFactory channelFactory,
       TikvBlockingStub blockingStub,
-      TikvStub asyncStub,
+      TikvFutureStub asyncStub,
       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,23 +114,23 @@ 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));
       }
       ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
 
       TikvBlockingStub tikvBlockingStub = TikvGrpc.newBlockingStub(channel);
-      TikvStub tikvAsyncStub = TikvGrpc.newStub(channel);
+      TikvGrpc.TikvFutureStub tikvAsyncStub = TikvGrpc.newFutureStub(channel);
 
       this.lockResolverClient =
           AbstractLockResolverClient.getInstance(
-              tikvStore.getVersion(),
               conf,
               region,
+              tikvStore,
               tikvBlockingStub,
               tikvAsyncStub,
               channelFactory,
@@ -169,7 +171,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<GetRequest> factory =
         () ->
             GetRequest.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(version)))
+                .setContext(makeContext(getResolvedLocks(version), this.storeType))
                 .setKey(key)
                 .setVersion(version)
                 .build();
@@ -214,7 +216,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<BatchGetRequest> request =
         () ->
             BatchGetRequest.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(version)))
+                .setContext(makeContext(getResolvedLocks(version), this.storeType))
                 .addAllKeys(keys)
                 .setVersion(version)
                 .build();
@@ -277,7 +279,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<ScanRequest> request =
           () ->
               ScanRequest.newBuilder()
-                  .setContext(region.getContext(getResolvedLocks(version)))
+                  .setContext(makeContext(getResolvedLocks(version), this.storeType))
                   .setStartKey(startKey)
                   .setVersion(version)
                   .setKeyOnly(keyOnly)
@@ -379,7 +381,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
           () ->
               getIsV4()
                   ? PrewriteRequest.newBuilder()
-                      .setContext(region.getContext())
+                      .setContext(makeContext(storeType))
                       .setStartVersion(startTs)
                       .setPrimaryLock(primaryLock)
                       .addAllMutations(mutations)
@@ -389,7 +391,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
                       .setTxnSize(16)
                       .build()
                   : PrewriteRequest.newBuilder()
-                      .setContext(region.getContext())
+                      .setContext(makeContext(storeType))
                       .setStartVersion(startTs)
                       .setPrimaryLock(primaryLock)
                       .addAllMutations(mutations)
@@ -469,7 +471,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<TxnHeartBeatRequest> factory =
           () ->
               TxnHeartBeatRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .setStartVersion(startTs)
                   .setPrimaryLock(primaryLock)
                   .setAdviseLockTtl(ttl)
@@ -527,7 +529,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
                 .setStartVersion(startTs)
                 .setCommitVersion(commitTs)
                 .addAllKeys(keys)
-                .setContext(region.getContext())
+                .setContext(makeContext(storeType))
                 .build();
     KVErrorHandler<CommitResponse> handler =
         new KVErrorHandler<>(
@@ -588,7 +590,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<Coprocessor.Request> reqToSend =
         () ->
             Coprocessor.Request.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(startTs)))
+                .setContext(makeContext(getResolvedLocks(startTs), this.storeType))
                 .setTp(REQ_TYPE_DAG.getValue())
                 .setStartTs(startTs)
                 .setData(req.toByteString())
@@ -711,7 +713,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     Supplier<Coprocessor.Request> reqToSend =
         () ->
             Coprocessor.Request.newBuilder()
-                .setContext(region.getContext(getResolvedLocks(startTs)))
+                .setContext(makeContext(getResolvedLocks(startTs), this.storeType))
                 // TODO: If no executors...?
                 .setTp(REQ_TYPE_DAG.getValue())
                 .setData(req.toByteString())
@@ -745,11 +747,11 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
    * @param splitKeys is the split points for a specific region.
    * @return a split region info.
    */
-  public List<TiRegion> splitRegion(Iterable<ByteString> splitKeys) {
+  public List<Metapb.Region> splitRegion(Iterable<ByteString> splitKeys) {
     Supplier<SplitRegionRequest> request =
         () ->
             SplitRegionRequest.newBuilder()
-                .setContext(region.getContext())
+                .setContext(makeContext(storeType))
                 .addAllSplitKeys(splitKeys)
                 .build();
 
@@ -780,18 +782,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
               region.getId(), resp.getRegionError().toString()));
     }
 
-    return resp.getRegionsList()
-        .stream()
-        .map(
-            region ->
-                new TiRegion(
-                    region,
-                    null,
-                    conf.getIsolationLevel(),
-                    conf.getCommandPriority(),
-                    conf.getKvMode(),
-                    conf.isReplicaRead()))
-        .collect(Collectors.toList());
+    return resp.getRegionsList();
   }
 
   // APIs for Raw Scan/Put/Get/Delete
@@ -801,9 +792,9 @@ 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();
-      KVErrorHandler<RawGetResponse> handler =
-          new KVErrorHandler<>(
+          () -> RawGetRequest.newBuilder().setContext(makeContext(storeType)).setKey(key).build();
+      RegionErrorHandler<RawGetResponse> handler =
+          new RegionErrorHandler<RawGetResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawGetResponse resp = callWithRetry(backOffer, TikvGrpc.getRawGetMethod(), factory, handler);
       return rawGetHelper(resp);
@@ -833,9 +824,12 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     try {
       Supplier<RawGetKeyTTLRequest> factory =
           () ->
-              RawGetKeyTTLRequest.newBuilder().setContext(region.getContext()).setKey(key).build();
-      KVErrorHandler<RawGetKeyTTLResponse> handler =
-          new KVErrorHandler<>(
+              RawGetKeyTTLRequest.newBuilder()
+                  .setContext(makeContext(storeType))
+                  .setKey(key)
+                  .build();
+      RegionErrorHandler<RawGetKeyTTLResponse> handler =
+          new RegionErrorHandler<RawGetKeyTTLResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawGetKeyTTLResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawGetKeyTTLMethod(), factory, handler);
@@ -863,15 +857,20 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     return resp.getTtl();
   }
 
-  public void rawDelete(BackOffer backOffer, ByteString key) {
+  public void rawDelete(BackOffer backOffer, ByteString key, boolean atomic) {
     Histogram.Timer requestTimer =
         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(makeContext(storeType))
+                  .setKey(key)
+                  .setForCas(atomic)
+                  .build();
 
-      KVErrorHandler<RawDeleteResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawDeleteResponse> handler =
+          new RegionErrorHandler<RawDeleteResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawDeleteResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawDeleteMethod(), factory, handler);
@@ -895,21 +894,23 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
     }
   }
 
-  public void rawPut(BackOffer backOffer, ByteString key, ByteString value, long ttl) {
+  public void rawPut(
+      BackOffer backOffer, ByteString key, ByteString value, long ttl, boolean atomic) {
     Histogram.Timer requestTimer =
         GRPC_RAW_REQUEST_LATENCY.labels("client_grpc_raw_put").startTimer();
     try {
       Supplier<RawPutRequest> factory =
           () ->
               RawPutRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .setKey(key)
                   .setValue(value)
                   .setTtl(ttl)
+                  .setForCas(atomic)
                   .build();
 
-      KVErrorHandler<RawPutResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawPutResponse> handler =
+          new RegionErrorHandler<RawPutResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawPutResponse resp = callWithRetry(backOffer, TikvGrpc.getRawPutMethod(), factory, handler);
       rawPutHelper(resp);
@@ -940,15 +941,15 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawCASRequest> factory =
           () ->
               RawCASRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .setKey(key)
                   .setValue(value)
                   .setPreviousNotExist(true)
                   .setTtl(ttl)
                   .build();
 
-      KVErrorHandler<RawCASResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawCASResponse> handler =
+          new RegionErrorHandler<RawCASResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawCASResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawCompareAndSwapMethod(), factory, handler);
@@ -961,7 +962,7 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
   private ByteString rawPutIfAbsentHelper(RawCASResponse resp) {
     if (resp == null) {
       this.regionManager.onRequestFail(region);
-      throw new TiClientInternalException("RawPutResponse failed without a cause");
+      throw new TiClientInternalException("RawCASResponse failed without a cause");
     }
     String error = resp.getError();
     if (!error.isEmpty()) {
@@ -986,11 +987,11 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchGetRequest> factory =
           () ->
               RawBatchGetRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .addAllKeys(keys)
                   .build();
-      KVErrorHandler<RawBatchGetResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawBatchGetResponse> handler =
+          new RegionErrorHandler<RawBatchGetResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawBatchGetResponse resp =
           callWithRetry(backoffer, TikvGrpc.getRawBatchGetMethod(), factory, handler);
@@ -1021,13 +1022,13 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchPutRequest> factory =
           () ->
               RawBatchPutRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .addAllPairs(kvPairs)
                   .setTtl(ttl)
                   .setForCas(atomic)
                   .build();
-      KVErrorHandler<RawBatchPutResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawBatchPutResponse> handler =
+          new RegionErrorHandler<RawBatchPutResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawBatchPutResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawBatchPutMethod(), factory, handler);
@@ -1073,12 +1074,12 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawBatchDeleteRequest> factory =
           () ->
               RawBatchDeleteRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .addAllKeys(keys)
                   .setForCas(atomic)
                   .build();
-      KVErrorHandler<RawBatchDeleteResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawBatchDeleteResponse> handler =
+          new RegionErrorHandler<RawBatchDeleteResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawBatchDeleteResponse resp =
           callWithRetry(backoffer, TikvGrpc.getRawBatchDeleteMethod(), factory, handler);
@@ -1118,14 +1119,14 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawScanRequest> factory =
           () ->
               RawScanRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .setStartKey(key)
                   .setKeyOnly(keyOnly)
                   .setLimit(limit)
                   .build();
 
-      KVErrorHandler<RawScanResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawScanResponse> handler =
+          new RegionErrorHandler<RawScanResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawScanResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawScanMethod(), factory, handler);
@@ -1164,13 +1165,13 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       Supplier<RawDeleteRangeRequest> factory =
           () ->
               RawDeleteRangeRequest.newBuilder()
-                  .setContext(region.getContext())
+                  .setContext(makeContext(storeType))
                   .setStartKey(startKey)
                   .setEndKey(endKey)
                   .build();
 
-      KVErrorHandler<RawDeleteRangeResponse> handler =
-          new KVErrorHandler<>(
+      RegionErrorHandler<RawDeleteRangeResponse> handler =
+          new RegionErrorHandler<RawDeleteRangeResponse>(
               regionManager, this, resp -> resp.hasRegionError() ? resp.getRegionError() : null);
       RawDeleteRangeResponse resp =
           callWithRetry(backOffer, TikvGrpc.getRawDeleteRangeMethod(), factory, handler);
@@ -1232,25 +1233,39 @@ 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);
-      TikvStub asyncStub = TikvGrpc.newStub(channel);
+      TikvBlockingStub blockingStub = null;
+      TikvFutureStub asyncStub = null;
+
+      if (conf.getEnableGrpcForward() && store.getProxyStore() != null && !store.isReachable()) {
+        addressStr = store.getProxyStore().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.newFutureStub(channel), header);
+      } else {
+        channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+        blockingStub = TikvGrpc.newBlockingStub(channel);
+        asyncStub = TikvGrpc.newFutureStub(channel);
+      }
 
       return new RegionStoreClient(
           conf,
           region,
-          store.getVersion(),
+          store,
           storeType,
           channelFactory,
           blockingStub,
@@ -1260,7 +1275,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);
     }
 
@@ -1268,19 +1284,39 @@ public class RegionStoreClient extends AbstractRegionStoreClient {
       return build(key, TiStoreType.TiKV);
     }
 
+    public synchronized RegionStoreClient build(ByteString key, BackOffer backOffer)
+        throws GrpcException {
+      return build(key, TiStoreType.TiKV, backOffer);
+    }
+
     public synchronized RegionStoreClient build(ByteString key, TiStoreType storeType)
         throws GrpcException {
-      Pair<TiRegion, Store> pair = regionManager.getRegionStorePairByKey(key, storeType);
+      return build(key, storeType, defaultBackOff());
+    }
+
+    public synchronized RegionStoreClient build(
+        ByteString key, TiStoreType storeType, BackOffer backOffer) throws GrpcException {
+      Pair<TiRegion, TiStore> pair =
+          regionManager.getRegionStorePairByKey(key, storeType, backOffer);
       return build(pair.first, pair.second, storeType);
     }
 
     public synchronized RegionStoreClient build(TiRegion region) throws GrpcException {
-      Store store = regionManager.getStoreById(region.getLeader().getStoreId());
+      return build(region, defaultBackOff());
+    }
+
+    public synchronized RegionStoreClient build(TiRegion region, BackOffer backOffer)
+        throws GrpcException {
+      TiStore store = regionManager.getStoreById(region.getLeader().getStoreId(), backOffer);
       return build(region, store, TiStoreType.TiKV);
     }
 
     public RegionManager getRegionManager() {
       return regionManager;
     }
+
+    private BackOffer defaultBackOff() {
+      return ConcreteBackOffer.newCustomBackOff(conf.getRawKVDefaultBackoffInMS());
+    }
   }
 }

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

@@ -0,0 +1,146 @@
+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.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.ReadOnlyPDClient;
+import org.tikv.common.util.ChannelFactory;
+import org.tikv.common.util.ConcreteBackOffer;
+import org.tikv.kvproto.Metapb;
+
+public class StoreHealthyChecker implements Runnable {
+  private static final Logger logger = LoggerFactory.getLogger(StoreHealthyChecker.class);
+  private static final long MAX_CHECK_STORE_TOMBSTONE_TICK = 60;
+  private final BlockingQueue<TiStore> taskQueue;
+  private final ChannelFactory channelFactory;
+  private final ReadOnlyPDClient pdClient;
+  private final RegionCache cache;
+  private long checkTombstoneTick;
+  private final long timeout;
+
+  public StoreHealthyChecker(
+      ChannelFactory channelFactory, ReadOnlyPDClient pdClient, RegionCache cache, long timeout) {
+    this.taskQueue = new LinkedBlockingQueue<>();
+    this.channelFactory = channelFactory;
+    this.pdClient = pdClient;
+    this.cache = cache;
+    this.checkTombstoneTick = 0;
+    this.timeout = timeout;
+  }
+
+  public boolean scheduleStoreHealthCheck(TiStore store) {
+    if (!this.taskQueue.add(store)) {
+      // add queue false, mark it reachable so that it can be put again.
+      return false;
+    }
+    return true;
+  }
+
+  private List<TiStore> getValidStores() {
+    List<TiStore> unhealthStore = new LinkedList<>();
+    while (!this.taskQueue.isEmpty()) {
+      try {
+        TiStore store = this.taskQueue.take();
+        if (!store.isValid()) {
+          continue;
+        }
+        unhealthStore.add(store);
+      } catch (Exception e) {
+        return unhealthStore;
+      }
+    }
+    return unhealthStore;
+  }
+
+  private boolean checkStoreHealth(TiStore store) {
+    String addressStr = store.getStore().getAddress();
+    try {
+      ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+      HealthGrpc.HealthBlockingStub stub =
+          HealthGrpc.newBlockingStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
+      HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
+      HealthCheckResponse resp = stub.check(req);
+      if (resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING) {
+        return true;
+      } else {
+        return false;
+      }
+    } catch (Exception e) {
+      return false;
+    }
+  }
+
+  private boolean checkStoreTombstone(TiStore store) {
+    try {
+      Metapb.Store newStore = pdClient.getStore(ConcreteBackOffer.newRawKVBackOff(), store.getId());
+      if (newStore != null && newStore.getState() == Metapb.StoreState.Tombstone) {
+        return true;
+      }
+    } catch (Exception e) {
+      return false;
+    }
+    return false;
+  }
+
+  @Override
+  public void run() {
+    checkTombstoneTick += 1;
+    boolean needCheckTombstoneStore = false;
+    if (checkTombstoneTick >= MAX_CHECK_STORE_TOMBSTONE_TICK) {
+      needCheckTombstoneStore = true;
+      checkTombstoneTick = 0;
+    }
+    List<TiStore> allStores = getValidStores();
+    List<TiStore> unreachableStore = new LinkedList<>();
+    for (TiStore store : allStores) {
+      if (needCheckTombstoneStore) {
+        if (checkStoreTombstone(store)) {
+          continue;
+        }
+      }
+
+      if (checkStoreHealth(store)) {
+        if (store.getProxyStore() != null) {
+          TiStore newStore = store.withProxy(null);
+          logger.warn(String.format("store [%s] recovers to be reachable", store.getAddress()));
+          if (cache.putStore(newStore.getId(), newStore)) {
+            this.taskQueue.add(newStore);
+            continue;
+          }
+        } else {
+          if (!store.isReachable()) {
+            logger.warn(String.format("store [%s] recovers to be reachable", store.getAddress()));
+            store.markReachable();
+          }
+        }
+      } else if (store.isReachable()) {
+        unreachableStore.add(store);
+        continue;
+      }
+      this.taskQueue.add(store);
+    }
+    if (!unreachableStore.isEmpty()) {
+      try {
+        Thread.sleep(timeout);
+      } catch (Exception e) {
+        this.taskQueue.addAll(unreachableStore);
+        return;
+      }
+      for (TiStore store : unreachableStore) {
+        if (!checkStoreHealth(store)) {
+          logger.warn(String.format("store [%s] is not reachable", store.getAddress()));
+          store.markUnreachable();
+        }
+        this.taskQueue.add(store);
+      }
+    }
+  }
+}

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

@@ -22,14 +22,15 @@ 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.tikv.common.TiConfiguration.KVMode;
-import org.tikv.common.codec.Codec.BytesCodec;
-import org.tikv.common.codec.CodecDataInput;
+import java.util.stream.Collectors;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.TiConfiguration;
 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 +40,28 @@ 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 IsolationLevel isolationLevel;
   private final Kvrpcpb.CommandPri commandPri;
+  private final TiConfiguration conf;
   private final Peer leader;
-  private int followerIdx = 0;
-  private final boolean isReplicaRead;
+  private final ReplicaSelector replicaSelector;
+  private final List<Peer> replicaList;
+  private int replicaIdx;
+  private final List<Peer> peers;
+  private final List<TiStore> stores;
 
   public TiRegion(
-      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,
-      IsolationLevel isolationLevel,
-      Kvrpcpb.CommandPri commandPri,
-      KVMode kvMode,
-      boolean isReplicaRead) {
-    Objects.requireNonNull(meta, "meta is null");
-    this.meta = decodeRegion(meta, kvMode == KVMode.RAW);
+      TiConfiguration conf, Region meta, Peer leader, List<Peer> peers, List<TiStore> stores) {
+    this.conf = Objects.requireNonNull(conf, "conf is null");
+    this.meta = Objects.requireNonNull(meta, "meta is null");
+    this.isolationLevel = conf.getIsolationLevel();
+    this.commandPri = conf.getCommandPriority();
+    this.peers = peers;
+    this.stores = stores;
+    this.replicaSelector = conf.getReplicaSelector();
     if (leader == null || leader.getId() == 0) {
       if (meta.getPeersCount() == 0) {
         throw new TiClientInternalException("Empty peer list for region " + meta.getId());
@@ -86,77 +71,60 @@ public class TiRegion implements Serializable {
     } else {
       this.leader = leader;
     }
-    if (isReplicaRead && meta.getPeersCount() > 0) {
-      // try to get first follower
-      try {
-        chooseRandomFollower();
-      } catch (Exception ignore) {
-        // ignore
-      }
-    }
-    this.isolationLevel = isolationLevel;
-    this.commandPri = commandPri;
-    this.isReplicaRead = isReplicaRead;
+
+    // init replicaList
+    replicaList =
+        replicaSelector
+            .select(new org.tikv.common.replica.Region(meta, this.leader, peers, stores))
+            .stream()
+            .map(org.tikv.common.replica.Store::getPeer)
+            .collect(Collectors.toList());
+    replicaIdx = 0;
   }
 
-  private Region decodeRegion(Region region, boolean isRawRegion) {
-    Region.Builder builder =
-        Region.newBuilder()
-            .setId(region.getId())
-            .setRegionEpoch(region.getRegionEpoch())
-            .addAllPeers(region.getPeersList());
-
-    if (region.getStartKey().isEmpty() || isRawRegion) {
-      builder.setStartKey(region.getStartKey());
-    } else {
-      byte[] decodedStartKey = BytesCodec.readBytes(new CodecDataInput(region.getStartKey()));
-      builder.setStartKey(ByteString.copyFrom(decodedStartKey));
-    }
-
-    if (region.getEndKey().isEmpty() || isRawRegion) {
-      builder.setEndKey(region.getEndKey());
-    } else {
-      byte[] decodedEndKey = BytesCodec.readBytes(new CodecDataInput(region.getEndKey()));
-      builder.setEndKey(ByteString.copyFrom(decodedEndKey));
-    }
-
-    return builder.build();
+  public TiConfiguration getConf() {
+    return conf;
   }
 
   public Peer getLeader() {
     return leader;
   }
 
-  public Peer getCurrentFollower() {
-    return meta.getPeers(followerIdx);
-  }
-
-  private boolean isValidFollower(Peer peer) {
-    return Metapb.PeerRole.valueOf(peer.getRole().getValueDescriptor()) == Metapb.PeerRole.Voter;
-  }
-
-  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;
+  public List<Peer> getFollowerList() {
+    List<Peer> peers = new ArrayList<>();
+    for (Peer peer : getMeta().getPeersList()) {
+      if (!peer.equals(this.leader)) {
+        if (peer.getRole().equals(Metapb.PeerRole.Voter)) {
+          peers.add(peer);
+        }
       }
     }
+    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 +145,35 @@ public class TiRegion implements Serializable {
     return Key.toRawKey(getEndKey());
   }
 
-  public Kvrpcpb.Context getContext() {
-    return getContext(java.util.Collections.emptySet());
+  public Kvrpcpb.Context getLeaderContext() {
+    return getContext(this.leader, java.util.Collections.emptySet(), false);
   }
 
-  public Kvrpcpb.Context getContext(Set<Long> resolvedLocks) {
+  public Kvrpcpb.Context getReplicaContext(Set<Long> resolvedLocks, TiStoreType storeType) {
+    Peer currentPeer = getCurrentReplica();
+    boolean replicaRead = !isLeader(currentPeer) && TiStoreType.TiKV.equals(storeType);
+    return getContext(currentPeer, resolvedLocks, replicaRead);
+  }
+
+  public Kvrpcpb.Context getReplicaContext(Peer currentPeer, Set<Long> resolvedLocks) {
+    return getContext(currentPeer, resolvedLocks, false);
+  }
+
+  public Kvrpcpb.Context getReplicaContext(Peer currentPeer) {
+    return getContext(currentPeer, java.util.Collections.emptySet(), false);
+  }
+
+  private Kvrpcpb.Context getContext(
+      Peer currentPeer, Set<Long> resolvedLocks, boolean replicaRead) {
+
     Kvrpcpb.Context.Builder builder = Kvrpcpb.Context.newBuilder();
-    builder.setIsolationLevel(this.isolationLevel);
-    builder.setPriority(this.commandPri);
-    if (isReplicaRead) {
-      builder
-          .setRegionId(meta.getId())
-          .setPeer(getCurrentFollower())
-          .setReplicaRead(true)
-          .setRegionEpoch(this.meta.getRegionEpoch());
-    } else {
-      builder
-          .setRegionId(meta.getId())
-          .setPeer(this.leader)
-          .setRegionEpoch(this.meta.getRegionEpoch());
-    }
+    builder
+        .setIsolationLevel(this.isolationLevel)
+        .setPriority(this.commandPri)
+        .setRegionId(meta.getId())
+        .setPeer(currentPeer)
+        .setReplicaRead(replicaRead)
+        .setRegionEpoch(this.meta.getRegionEpoch());
     builder.addAllResolvedLocks(resolvedLocks);
     return builder.build();
   }
@@ -218,7 +195,7 @@ 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.conf, this.meta, p, peers, this.stores);
       }
     }
     return null;
@@ -302,6 +279,18 @@ public class TiRegion implements Serializable {
       this.ver = ver;
     }
 
+    public long getId() {
+      return id;
+    }
+
+    public long getConfVer() {
+      return confVer;
+    }
+
+    public long getVer() {
+      return ver;
+    }
+
     @Override
     public boolean equals(Object other) {
       if (this == other) {

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

@@ -0,0 +1,91 @@
+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 final Metapb.Store proxyStore;
+  private final AtomicBoolean reachable;
+  private final AtomicBoolean valid;
+
+  public TiStore(Metapb.Store store) {
+    this.store = store;
+    this.reachable = new AtomicBoolean(true);
+    this.valid = new AtomicBoolean(true);
+    this.proxyStore = null;
+  }
+
+  private TiStore(Metapb.Store store, Metapb.Store proxyStore) {
+    this.store = store;
+    if (proxyStore != null) {
+      this.reachable = new AtomicBoolean(false);
+    } else {
+      this.reachable = new AtomicBoolean(true);
+    }
+    this.valid = new AtomicBoolean(true);
+    this.proxyStore = proxyStore;
+  }
+
+  @java.lang.Override
+  public boolean equals(final java.lang.Object obj) {
+    if (obj == this) {
+      return true;
+    }
+    if (!(obj instanceof TiStore)) {
+      return super.equals(obj);
+    }
+    TiStore other = (TiStore) obj;
+    if (!this.store.equals(other.store)) {
+      return false;
+    }
+
+    if (proxyStore == null && other.proxyStore == null) {
+      return true;
+    }
+    if (proxyStore != null && other.proxyStore != null) {
+      return proxyStore.equals(other.proxyStore);
+    }
+    return false;
+  }
+
+  public TiStore withProxy(Metapb.Store proxyStore) {
+    return new TiStore(this.store, proxyStore);
+  }
+
+  public void markUnreachable() {
+    this.reachable.set(false);
+  }
+
+  public void markReachable() {
+    this.reachable.set(true);
+  }
+
+  public boolean isReachable() {
+    return this.reachable.get();
+  }
+
+  public boolean isValid() {
+    return this.valid.get();
+  }
+
+  public void markInvalid() {
+    this.valid.set(false);
+  }
+
+  public Metapb.Store getStore() {
+    return this.store;
+  }
+
+  public String getAddress() {
+    return this.store.getAddress();
+  }
+
+  public Metapb.Store getProxyStore() {
+    return this.proxyStore;
+  }
+
+  public long getId() {
+    return this.store.getId();
+  }
+}

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

@@ -0,0 +1,36 @@
+/*
+ * 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;
+
+public class FollowerReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Store> select(Region region) {
+    Store[] stores = region.getStores();
+    Store leader = region.getLeader();
+    List<Store> list = new ArrayList<>(stores.length);
+    for (Store store : stores) {
+      if (!store.isLearner() && !leader.equals(store)) {
+        list.add(store);
+      }
+    }
+    Collections.shuffle(list);
+    return list;
+  }
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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;
+
+public class LeaderFollowerReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Store> select(Region region) {
+    Store[] stores = region.getStores();
+    Store leader = region.getLeader();
+    List<Store> list = new ArrayList<>(stores.length);
+    for (Store store : stores) {
+      if (!store.isLearner() && !leader.equals(store)) {
+        list.add(store);
+      }
+    }
+    Collections.shuffle(list);
+    list.add(leader);
+    return list;
+  }
+}

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

@@ -0,0 +1,28 @@
+/*
+ * 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;
+
+public class LeaderReplicaSelector implements ReplicaSelector {
+  @Override
+  public List<Store> select(Region region) {
+    List<Store> list = new ArrayList<>(1);
+    list.add(region.getLeader());
+    return list;
+  }
+}

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

@@ -0,0 +1,57 @@
+package org.tikv.common.replica;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+
+import java.util.Iterator;
+import java.util.List;
+import org.tikv.common.region.TiStore;
+import org.tikv.kvproto.Metapb;
+
+public class Region {
+  private final Metapb.Region region;
+  private final Store[] stores;
+  private Store leaderStore;
+
+  public Region(
+      final Metapb.Region region,
+      final Metapb.Peer leader,
+      final List<Metapb.Peer> peers,
+      final List<TiStore> stores) {
+    this.region = region;
+    this.stores = new Store[stores.size()];
+    Iterator<Metapb.Peer> peer = peers.iterator();
+    Iterator<TiStore> store = stores.iterator();
+    for (int idx = 0; idx < peers.size(); idx++) {
+      Metapb.Peer currentPeer = peer.next();
+      boolean isLeader = currentPeer.equals(leader);
+      this.stores[idx] = new Store(currentPeer, store.next().getStore(), isLeader);
+      if (isLeader) {
+        leaderStore = this.stores[idx];
+      }
+    }
+  }
+
+  public Store[] getStores() {
+    return stores;
+  }
+
+  public Store getLeader() {
+    return leaderStore;
+  }
+
+  public long getId() {
+    return region.getId();
+  }
+
+  public byte[] getStartKey() {
+    return region.getStartKey().toByteArray();
+  }
+
+  public byte[] getEndKey() {
+    return region.getEndKey().toByteArray();
+  }
+
+  public String toString() {
+    return toStringHelper(this).add("region", region).add("stores", stores).toString();
+  }
+}

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

@@ -0,0 +1,27 @@
+/*
+ * 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;
+
+public interface ReplicaSelector extends Serializable {
+  ReplicaSelector LEADER = new LeaderReplicaSelector();
+  ReplicaSelector FOLLOWER = new FollowerReplicaSelector();
+  ReplicaSelector LEADER_AND_FOLLOWER = new LeaderFollowerReplicaSelector();
+
+  List<Store> select(Region region);
+}

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

@@ -0,0 +1,111 @@
+package org.tikv.common.replica;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+
+import java.util.List;
+import org.tikv.kvproto.Metapb;
+
+public class Store {
+  public static class Label {
+    private final org.tikv.kvproto.Metapb.StoreLabel label;
+
+    Label(org.tikv.kvproto.Metapb.StoreLabel label) {
+      this.label = label;
+    }
+
+    public String getKey() {
+      return label.getKey();
+    }
+
+    public String getValue() {
+      return label.getValue();
+    }
+  }
+
+  public enum State {
+    Unknown,
+    Up,
+    Offline,
+    Tombstone
+  }
+
+  private static final Label[] EMPTY_LABELS = new Label[0];
+  private Label[] labels;
+  private final Metapb.Peer peer;
+  private final Metapb.Store store;
+  private final boolean isLeader;
+
+  Store(
+      final org.tikv.kvproto.Metapb.Peer peer,
+      final org.tikv.kvproto.Metapb.Store store,
+      boolean isLeader) {
+    this.peer = peer;
+    this.store = store;
+    this.isLeader = isLeader;
+  }
+
+  public Metapb.Peer getPeer() {
+    return peer;
+  }
+
+  public Label[] getLabels() {
+    if (labels == null) {
+      List<Metapb.StoreLabel> labelList = store.getLabelsList();
+      if (labelList.isEmpty()) {
+        labels = EMPTY_LABELS;
+      } else {
+        labels = labelList.stream().map(Label::new).toArray(Label[]::new);
+      }
+    }
+    return labels;
+  }
+
+  public boolean isLearner() {
+    return peer.getRole() == Metapb.PeerRole.Learner;
+  }
+
+  public boolean isLeader() {
+    return isLeader;
+  }
+
+  public boolean isFollower() {
+    return peer.getRole() == Metapb.PeerRole.Voter && !isLeader;
+  }
+
+  public long getId() {
+    return store.getId();
+  }
+
+  public String getAddress() {
+    return store.getAddress();
+  }
+
+  public String getVersion() {
+    return store.getVersion();
+  }
+
+  public State getState() {
+    switch (store.getState()) {
+      case Up:
+        return State.Up;
+      case Offline:
+        return State.Offline;
+      case Tombstone:
+        return State.Tombstone;
+      default:
+        return State.Unknown;
+    }
+  }
+
+  public boolean equals(Object o) {
+    if (!(o instanceof Store)) {
+      return false;
+    }
+    Store other = (Store) o;
+    return this.peer.equals(other.peer);
+  }
+
+  public String toString() {
+    return toStringHelper(this).add("peer", peer).add("store", store).toString();
+  }
+}

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

@@ -1,7 +1,6 @@
 package org.tikv.common.util;
 
 import java.util.concurrent.ThreadLocalRandom;
-import org.tikv.common.exception.GrpcException;
 
 public class BackOffFunction {
   private final int base;
@@ -25,7 +24,7 @@ public class BackOffFunction {
    * Do back off in exponential with optional jitters according to different back off strategies.
    * See http://www.awsarchitectureblog.com/2015/03/backoff.html
    */
-  long doBackOff(long maxSleepMs) {
+  long getSleepMs(long maxSleepMs) {
     long sleep = 0;
     long v = expo(base, cap, attempts);
     switch (strategy) {
@@ -47,11 +46,6 @@ public class BackOffFunction {
       sleep = maxSleepMs;
     }
 
-    try {
-      Thread.sleep(sleep);
-    } catch (InterruptedException e) {
-      throw new GrpcException(e);
-    }
     attempts++;
     lastSleep = sleep;
     return lastSleep;
@@ -69,6 +63,7 @@ public class BackOffFunction {
     BoRegionMiss,
     BoUpdateLeader,
     BoServerBusy,
-    BoTxnNotFound
+    BoTxnNotFound,
+    BoCheckTimeout
   }
 }

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

@@ -17,25 +17,17 @@
 
 package org.tikv.common.util;
 
+import org.tikv.common.log.SlowLog;
+
 public interface BackOffer {
   // Back off types.
   int seconds = 1000;
-  int COP_BUILD_TASK_MAX_BACKOFF = 5 * seconds;
   int TSO_MAX_BACKOFF = 5 * seconds;
   int SCANNER_NEXT_MAX_BACKOFF = 40 * seconds;
   int BATCH_GET_MAX_BACKOFF = 40 * seconds;
   int COP_NEXT_MAX_BACKOFF = 40 * seconds;
   int GET_MAX_BACKOFF = 40 * seconds;
-  int PREWRITE_MAX_BACKOFF = 20 * seconds;
-  int CLEANUP_MAX_BACKOFF = 20 * seconds;
-  int GC_ONE_REGION_MAX_BACKOFF = 20 * seconds;
-  int GC_RESOLVE_LOCK_MAX_BACKOFF = 100 * seconds;
-  int GC_DELETE_RANGE_MAX_BACKOFF = 100 * seconds;
-
   int RAWKV_MAX_BACKOFF = 20 * seconds;
-
-  int SPLIT_REGION_BACKOFF = 20 * seconds;
-  int BATCH_COMMIT_BACKOFF = 10 * seconds;
   int PD_INFO_BACKOFF = 5 * seconds;
 
   /**
@@ -44,6 +36,16 @@ public interface BackOffer {
    */
   void doBackOff(BackOffFunction.BackOffFuncType funcType, Exception err);
 
+  /** check if deadline exceeded. */
+  void checkTimeout();
+
+  /**
+   * canRetryAfterSleep sleeps a while base on the BackOffType and records the error message. Will
+   * stop until max back off time exceeded and throw an exception to the caller. It will return
+   * false if the total sleep time has exceed some limit condition.
+   */
+  boolean canRetryAfterSleep(BackOffFunction.BackOffFuncType funcType);
+
   /**
    * BackoffWithMaxSleep sleeps a while base on the backoffType and records the error message and
    * never sleep more than maxSleepMs for each sleep.
@@ -62,4 +64,6 @@ public interface BackOffer {
     // DecorrJitter increases the maximum jitter based on the last random value.
     DecorrJitter
   }
+
+  SlowLog getSlowLog();
 }

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/ChannelFactory.java

@@ -25,10 +25,12 @@ import org.tikv.common.pd.PDUtils;
 
 public class ChannelFactory implements AutoCloseable {
   private final int maxFrameSize;
+  private final int idleTimeout;
   private final ConcurrentHashMap<String, ManagedChannel> connPool = new ConcurrentHashMap<>();
 
-  public ChannelFactory(int maxFrameSize) {
+  public ChannelFactory(int maxFrameSize, int idleTimeout) {
     this.maxFrameSize = maxFrameSize;
+    this.idleTimeout = idleTimeout;
   }
 
   public ManagedChannel getChannel(String addressStr, HostMapping hostMapping) {
@@ -51,8 +53,8 @@ public class ChannelFactory implements AutoCloseable {
           // So a coarse grain lock is ok here
           return ManagedChannelBuilder.forAddress(mappedAddr.getHost(), mappedAddr.getPort())
               .maxInboundMessageSize(maxFrameSize)
-              .usePlaintext(true)
-              .idleTimeout(60, TimeUnit.SECONDS)
+              .usePlaintext()
+              .idleTimeout(idleTimeout, TimeUnit.SECONDS)
               .build();
         });
   }

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

@@ -19,6 +19,7 @@ import com.google.protobuf.ByteString;
 import java.util.*;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorCompletionService;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import org.tikv.common.exception.TiKVException;
@@ -121,7 +122,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);
   }
 
   /**
@@ -170,10 +171,14 @@ public class ClientUtils {
       ExecutorCompletionService<List<T>> completionService,
       Queue<List<T>> taskQueue,
       List<T> batches,
-      int backOff) {
+      long backOff) {
     try {
       for (int i = 0; i < batches.size(); i++) {
-        List<T> task = completionService.take().get(backOff, TimeUnit.MILLISECONDS);
+        Future<List<T>> future = completionService.poll(backOff, TimeUnit.MILLISECONDS);
+        if (future == null) {
+          throw new TiKVException("TimeOut Exceeded for current operation.");
+        }
+        List<T> task = future.get();
         if (!task.isEmpty()) {
           taskQueue.offer(task);
         }
@@ -181,8 +186,6 @@ public class ClientUtils {
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new TiKVException("Current thread interrupted.", e);
-    } catch (TimeoutException e) {
-      throw new TiKVException("TimeOut Exceeded for current operation. ", e);
     } catch (ExecutionException e) {
       throw new TiKVException("Execution exception met.", e);
     }
@@ -192,11 +195,16 @@ public class ClientUtils {
       ExecutorCompletionService<Pair<List<T>, List<U>>> completionService,
       Queue<List<T>> taskQueue,
       List<T> batches,
-      int backOff) {
+      long backOff) {
     try {
       List<U> result = new ArrayList<>();
       for (int i = 0; i < batches.size(); i++) {
-        Pair<List<T>, List<U>> task = completionService.take().get(backOff, TimeUnit.MILLISECONDS);
+        Future<Pair<List<T>, List<U>>> future =
+            completionService.poll(backOff, TimeUnit.MILLISECONDS);
+        if (future == null) {
+          throw new TiKVException("TimeOut Exceeded for current operation.");
+        }
+        Pair<List<T>, List<U>> task = future.get();
         if (!task.first.isEmpty()) {
           taskQueue.offer(task.first);
         } else {
@@ -207,8 +215,6 @@ public class ClientUtils {
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
       throw new TiKVException("Current thread interrupted.", e);
-    } catch (TimeoutException e) {
-      throw new TiKVException("TimeOut Exceeded for current operation. ", e);
     } catch (ExecutionException e) {
       throw new TiKVException("Execution exception met.", e);
     }

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

@@ -17,14 +17,22 @@
 
 package org.tikv.common.util;
 
+import static org.tikv.common.ConfigUtils.TIKV_BO_REGION_MISS_BASE_IN_MS;
+
 import com.google.common.base.Preconditions;
+import io.prometheus.client.Histogram;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.tikv.common.TiConfiguration;
 import org.tikv.common.exception.GrpcException;
+import org.tikv.common.exception.TiKVException;
+import org.tikv.common.log.SlowLog;
+import org.tikv.common.log.SlowLogEmptyImpl;
+import org.tikv.common.log.SlowLogSpan;
 
 public class ConcreteBackOffer implements BackOffer {
   private static final Logger logger = LoggerFactory.getLogger(ConcreteBackOffer.class);
@@ -32,12 +40,26 @@ public class ConcreteBackOffer implements BackOffer {
   private final Map<BackOffFunction.BackOffFuncType, BackOffFunction> backOffFunctionMap;
   private final List<Exception> errors;
   private int totalSleep;
+  private final long deadline;
+  private final SlowLog slowLog;
 
-  private ConcreteBackOffer(int maxSleep) {
+  public static final Histogram BACKOFF_DURATION =
+      Histogram.build()
+          .name("client_java_backoff_duration")
+          .help("backoff duration.")
+          .labelNames("type")
+          .register();
+
+  private ConcreteBackOffer(int maxSleep, long deadline, SlowLog slowLog) {
+    Preconditions.checkArgument(
+        maxSleep == 0 || deadline == 0, "Max sleep time should be 0 or Deadline should be 0.");
     Preconditions.checkArgument(maxSleep >= 0, "Max sleep time cannot be less than 0.");
+    Preconditions.checkArgument(deadline >= 0, "Deadline cannot be less than 0.");
     this.maxSleep = maxSleep;
     this.errors = new ArrayList<>();
     this.backOffFunctionMap = new HashMap<>();
+    this.deadline = deadline;
+    this.slowLog = slowLog;
   }
 
   private ConcreteBackOffer(ConcreteBackOffer source) {
@@ -45,34 +67,41 @@ public class ConcreteBackOffer implements BackOffer {
     this.totalSleep = source.totalSleep;
     this.errors = source.errors;
     this.backOffFunctionMap = source.backOffFunctionMap;
+    this.deadline = source.deadline;
+    this.slowLog = source.slowLog;
+  }
+
+  public static ConcreteBackOffer newDeadlineBackOff(int timeoutInMs, SlowLog slowLog) {
+    long deadline = System.currentTimeMillis() + timeoutInMs;
+    return new ConcreteBackOffer(0, deadline, slowLog);
   }
 
   public static ConcreteBackOffer newCustomBackOff(int maxSleep) {
-    return new ConcreteBackOffer(maxSleep);
+    return new ConcreteBackOffer(maxSleep, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newScannerNextMaxBackOff() {
-    return new ConcreteBackOffer(SCANNER_NEXT_MAX_BACKOFF);
+    return new ConcreteBackOffer(SCANNER_NEXT_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newBatchGetMaxBackOff() {
-    return new ConcreteBackOffer(BATCH_GET_MAX_BACKOFF);
+    return new ConcreteBackOffer(BATCH_GET_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newCopNextMaxBackOff() {
-    return new ConcreteBackOffer(COP_NEXT_MAX_BACKOFF);
+    return new ConcreteBackOffer(COP_NEXT_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newGetBackOff() {
-    return new ConcreteBackOffer(GET_MAX_BACKOFF);
+    return new ConcreteBackOffer(GET_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newRawKVBackOff() {
-    return new ConcreteBackOffer(RAWKV_MAX_BACKOFF);
+    return new ConcreteBackOffer(RAWKV_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer newTsoBackOff() {
-    return new ConcreteBackOffer(TSO_MAX_BACKOFF);
+    return new ConcreteBackOffer(TSO_MAX_BACKOFF, 0, SlowLogEmptyImpl.INSTANCE);
   }
 
   public static ConcreteBackOffer create(BackOffer source) {
@@ -96,20 +125,27 @@ public class ConcreteBackOffer implements BackOffer {
         backOffFunction = BackOffFunction.create(2000, 10000, BackOffStrategy.EqualJitter);
         break;
       case BoRegionMiss:
-        backOffFunction = BackOffFunction.create(100, 500, BackOffStrategy.NoJitter);
+        backOffFunction =
+            BackOffFunction.create(
+                TiConfiguration.getInt(TIKV_BO_REGION_MISS_BASE_IN_MS),
+                500,
+                BackOffStrategy.NoJitter);
         break;
       case BoTxnLock:
         backOffFunction = BackOffFunction.create(200, 3000, BackOffStrategy.EqualJitter);
         break;
       case BoPDRPC:
-        backOffFunction = BackOffFunction.create(500, 3000, BackOffStrategy.EqualJitter);
+        backOffFunction = BackOffFunction.create(100, 600, BackOffStrategy.EqualJitter);
         break;
       case BoTiKVRPC:
-        backOffFunction = BackOffFunction.create(100, 2000, BackOffStrategy.EqualJitter);
+        backOffFunction = BackOffFunction.create(10, 400, BackOffStrategy.EqualJitter);
         break;
       case BoTxnNotFound:
         backOffFunction = BackOffFunction.create(2, 500, BackOffStrategy.NoJitter);
         break;
+      case BoCheckTimeout:
+        backOffFunction = BackOffFunction.create(0, 0, BackOffStrategy.NoJitter);
+        break;
     }
     return backOffFunction;
   }
@@ -120,32 +156,79 @@ public class ConcreteBackOffer implements BackOffer {
   }
 
   @Override
-  public void doBackOffWithMaxSleep(
-      BackOffFunction.BackOffFuncType funcType, long maxSleepMs, Exception err) {
+  public void checkTimeout() {
+    if (!canRetryAfterSleep(BackOffFunction.BackOffFuncType.BoCheckTimeout)) {
+      logThrowError(new TiKVException("Request Timeout"));
+    }
+  }
+
+  @Override
+  public boolean canRetryAfterSleep(BackOffFunction.BackOffFuncType funcType) {
+    return canRetryAfterSleep(funcType, -1);
+  }
+
+  public boolean canRetryAfterSleep(BackOffFunction.BackOffFuncType funcType, long maxSleepMs) {
+    SlowLogSpan slowLogSpan = getSlowLog().start("backoff " + funcType.name());
+    Histogram.Timer backOffTimer = BACKOFF_DURATION.labels(funcType.name()).startTimer();
     BackOffFunction backOffFunction =
         backOffFunctionMap.computeIfAbsent(funcType, this::createBackOffFunc);
 
-    // Back off will be done here
-    totalSleep += backOffFunction.doBackOff(maxSleepMs);
+    // Back off will not be done here
+    long sleep = backOffFunction.getSleepMs(maxSleepMs);
+    totalSleep += sleep;
+    // Check deadline
+    if (deadline > 0) {
+      long currentMs = System.currentTimeMillis();
+      if (currentMs + sleep >= deadline) {
+        logger.warn(String.format("Deadline %d is exceeded, errors:", deadline));
+        return false;
+      }
+    }
+
+    try {
+      Thread.sleep(sleep);
+    } catch (InterruptedException e) {
+      throw new GrpcException(e);
+    } finally {
+      slowLogSpan.end();
+      backOffTimer.observeDuration();
+    }
+    if (maxSleep > 0 && totalSleep >= maxSleep) {
+      logger.warn(String.format("BackOffer.maxSleep %dms is exceeded, errors:", maxSleep));
+      return false;
+    }
+    return true;
+  }
+
+  @Override
+  public void doBackOffWithMaxSleep(
+      BackOffFunction.BackOffFuncType funcType, long maxSleepMs, Exception err) {
     logger.debug(
         String.format(
             "%s, retry later(totalSleep %dms, maxSleep %dms)",
             err.getMessage(), totalSleep, maxSleep));
     errors.add(err);
-    if (maxSleep > 0 && totalSleep >= maxSleep) {
-      StringBuilder errMsg =
-          new StringBuilder(
-              String.format("BackOffer.maxSleep %dms is exceeded, errors:", maxSleep));
-      for (int i = 0; i < errors.size(); i++) {
-        Exception curErr = errors.get(i);
-        // Print only last 3 errors for non-DEBUG log levels.
-        if (logger.isDebugEnabled() || i >= errors.size() - 3) {
-          errMsg.append("\n").append(i).append(".").append(curErr.toString());
-        }
-      }
-      logger.warn(errMsg.toString());
-      // Use the last backoff type to generate an exception
-      throw new GrpcException("retry is exhausted.", err);
+    if (!canRetryAfterSleep(funcType, maxSleepMs)) {
+      logThrowError(err);
     }
   }
+
+  private void logThrowError(Exception err) {
+    StringBuilder errMsg = new StringBuilder();
+    for (int i = 0; i < errors.size(); i++) {
+      Exception curErr = errors.get(i);
+      // Print only last 3 errors for non-DEBUG log levels.
+      if (logger.isDebugEnabled() || i >= errors.size() - 3) {
+        errMsg.append("\n").append(i).append(".").append(curErr.toString());
+      }
+    }
+    logger.warn(errMsg.toString());
+    // Use the last backoff type to generate an exception
+    throw new GrpcException("retry is exhausted.", err);
+  }
+
+  @Override
+  public SlowLog getSlowLog() {
+    return slowLog;
+  }
 }

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 =
-        new TLongObjectHashMap<>();
-    Map<Pair<TiRegion, Metapb.Store>, TLongArrayList> result = new HashMap<>();
+    TLongObjectHashMap<Pair<TiRegion, TiStore>> idToRegionStorePair = new TLongObjectHashMap<>();
+    Map<Pair<TiRegion, TiStore>, 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(
-        TiRegion region, Metapb.Store store, List<KeyRange> ranges) {
+    public static RegionTask newInstance(TiRegion region, TiStore 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

@@ -27,8 +27,12 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.TiConfiguration;
 import org.tikv.common.TiSession;
+import org.tikv.common.codec.KeyUtils;
 import org.tikv.common.exception.TiKVException;
 import org.tikv.common.key.Key;
+import org.tikv.common.log.SlowLog;
+import org.tikv.common.log.SlowLogEmptyImpl;
+import org.tikv.common.log.SlowLogImpl;
 import org.tikv.common.operation.iterator.RawScanIterator;
 import org.tikv.common.region.RegionStoreClient;
 import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
@@ -36,7 +40,7 @@ import org.tikv.common.region.TiRegion;
 import org.tikv.common.util.*;
 import org.tikv.kvproto.Kvrpcpb.KvPair;
 
-public class RawKVClient implements AutoCloseable {
+public class RawKVClient implements RawKVClientBase {
   private final RegionStoreClientBuilder clientBuilder;
   private final TiConfiguration conf;
   private final ExecutorService batchGetThreadPool;
@@ -46,15 +50,6 @@ public class RawKVClient implements AutoCloseable {
   private final ExecutorService deleteRangeThreadPool;
   private static final Logger logger = LoggerFactory.getLogger(RawKVClient.class);
 
-  // https://www.github.com/pingcap/tidb/blob/master/store/tikv/rawkv.go
-  private static final int MAX_RAW_SCAN_LIMIT = 10240;
-  private static final int MAX_RAW_BATCH_LIMIT = 1024;
-  private static final int RAW_BATCH_PUT_SIZE = 1024 * 1024; // 1 MB
-  private static final int RAW_BATCH_GET_SIZE = 16 * 1024; // 16 K
-  private static final int RAW_BATCH_DELETE_SIZE = 16 * 1024; // 16 K
-  private static final int RAW_BATCH_SCAN_SIZE = 16;
-  private static final int RAW_BATCH_PAIR_COUNT = 512;
-
   public static final Histogram RAW_REQUEST_LATENCY =
       Histogram.build()
           .name("client_java_raw_requests_latency")
@@ -94,124 +89,115 @@ public class RawKVClient implements AutoCloseable {
   @Override
   public void close() {}
 
-  /**
-   * Put a raw key-value pair to TiKV
-   *
-   * @param key raw key
-   * @param value raw value
-   */
+  @Override
   public void put(ByteString key, ByteString value) {
     put(key, value, 0);
   }
 
-  /**
-   * Put a raw key-value pair to TiKV
-   *
-   * @param key raw key
-   * @param value raw value
-   * @param ttl the ttl of the key (in seconds), 0 means the key will never be outdated
-   */
+  @Override
   public void put(ByteString key, ByteString value, long ttl) {
+    put(key, value, ttl, false);
+  }
+
+  @Override
+  public void putAtomic(ByteString key, ByteString value, long ttl) {
+    put(key, value, ttl, true);
+  }
+
+  private void put(ByteString key, ByteString value, long ttl, boolean atomic) {
     String label = "client_raw_put";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVWriteSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "put");
+                put("key", KeyUtils.formatBytesUTF8(key));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVWriteTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
       while (true) {
-        RegionStoreClient client = clientBuilder.build(key);
-        try {
-          client.rawPut(backOffer, key, value, ttl);
+        try (RegionStoreClient client = clientBuilder.build(key, backOffer)) {
+          slowLog.addProperty("region", client.getRegion().toString());
+          client.rawPut(backOffer, key, value, ttl, atomic);
           RAW_REQUEST_SUCCESS.labels(label).inc();
           return;
         } catch (final TiKVException e) {
           backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+          logger.warn("Retry for put error", e);
         }
       }
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Put a key-value pair if it does not exist. This API is atomic.
-   *
-   * @param key key
-   * @param value value
-   * @return a ByteString. returns ByteString.EMPTY if the value is written successfully. returns
-   *     the previous key if the value already exists, and does not write to TiKV.
-   */
+  @Override
   public ByteString putIfAbsent(ByteString key, ByteString value) {
     return putIfAbsent(key, value, 0L);
   }
 
-  /**
-   * Put a key-value pair with TTL if it does not exist. This API is atomic.
-   *
-   * @param key key
-   * @param value value
-   * @param ttl TTL of key (in seconds), 0 means the key will never be outdated.
-   * @return a ByteString. returns ByteString.EMPTY if the value is written successfully. returns
-   *     the previous key if the value already exists, and does not write to TiKV.
-   */
+  @Override
   public ByteString putIfAbsent(ByteString key, ByteString value, long ttl) {
     String label = "client_raw_put_if_absent";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVWriteSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "putIfAbsent");
+                put("key", KeyUtils.formatBytesUTF8(key));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVWriteTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
       while (true) {
-        RegionStoreClient client = clientBuilder.build(key);
-        try {
+        try (RegionStoreClient client = clientBuilder.build(key, backOffer)) {
+          slowLog.addProperty("region", client.getRegion().toString());
           ByteString result = client.rawPutIfAbsent(backOffer, key, value, ttl);
           RAW_REQUEST_SUCCESS.labels(label).inc();
           return result;
         } catch (final TiKVException e) {
           backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+          logger.warn("Retry for putIfAbsent error", e);
         }
       }
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Put a set of raw key-value pair to TiKV, this API does not ensure the operation is atomic.
-   *
-   * @param kvPairs kvPairs
-   */
+  @Override
   public void batchPut(Map<ByteString, ByteString> kvPairs) {
     batchPut(kvPairs, 0);
   }
 
-  /**
-   * Put a set of raw key-value pair to TiKV, this API does not ensure the operation is atomic.
-   *
-   * @param kvPairs kvPairs
-   * @param ttl the TTL of keys to be put (in seconds), 0 means the keys will never be outdated
-   */
+  @Override
   public void batchPut(Map<ByteString, ByteString> kvPairs, long ttl) {
     batchPut(kvPairs, ttl, false);
   }
 
-  /**
-   * Put a set of raw key-value pair to TiKV, this API is atomic
-   *
-   * @param kvPairs kvPairs
-   */
+  @Override
   public void batchPutAtomic(Map<ByteString, ByteString> kvPairs) {
     batchPutAtomic(kvPairs, 0);
   }
 
-  /**
-   * Put a set of raw key-value pair to TiKV, this API is atomic.
-   *
-   * @param kvPairs kvPairs
-   * @param ttl the TTL of keys to be put (in seconds), 0 means the keys will never be outdated
-   */
+  @Override
   public void batchPutAtomic(Map<ByteString, ByteString> kvPairs, long ttl) {
     batchPut(kvPairs, ttl, true);
   }
@@ -219,82 +205,105 @@ public class RawKVClient implements AutoCloseable {
   private void batchPut(Map<ByteString, ByteString> kvPairs, long ttl, boolean atomic) {
     String label = "client_raw_batch_put";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVBatchWriteSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "batchPut");
+                put("keySize", String.valueOf(kvPairs.size()));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVBatchWriteTimeoutInMS(), slowLog);
     try {
-      doSendBatchPut(ConcreteBackOffer.newRawKVBackOff(), kvPairs, ttl, atomic);
+      long deadline = System.currentTimeMillis() + conf.getRawKVBatchWriteTimeoutInMS();
+      doSendBatchPut(backOffer, kvPairs, ttl, atomic, deadline);
       RAW_REQUEST_SUCCESS.labels(label).inc();
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Get a raw key-value pair from TiKV if key exists
-   *
-   * @param key raw key
-   * @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
-   */
+  @Override
   public ByteString get(ByteString key) {
     String label = "client_raw_get";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVReadSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "get");
+                put("key", KeyUtils.formatBytesUTF8(key));
+              }
+            });
+
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVReadTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
       while (true) {
-        RegionStoreClient client = clientBuilder.build(key);
-        try {
-          ByteString result = client.rawGet(defaultBackOff(), key);
+        try (RegionStoreClient client = clientBuilder.build(key, backOffer)) {
+          slowLog.addProperty("region", client.getRegion().toString());
+          ByteString result = client.rawGet(backOffer, key);
           RAW_REQUEST_SUCCESS.labels(label).inc();
           return result;
         } catch (final TiKVException e) {
           backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+          logger.warn("Retry for get error", e);
         }
       }
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Get a list of raw key-value pair from TiKV if key exists
-   *
-   * @param keys list of raw key
-   * @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
-   */
+  @Override
   public List<KvPair> batchGet(List<ByteString> keys) {
     String label = "client_raw_batch_get";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVBatchReadSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "batchGet");
+                put("keySize", String.valueOf(keys.size()));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVBatchReadTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
-      List<KvPair> result = doSendBatchGet(backOffer, keys);
+      long deadline = System.currentTimeMillis() + conf.getRawKVBatchReadTimeoutInMS();
+      List<KvPair> result = doSendBatchGet(backOffer, keys, deadline);
       RAW_REQUEST_SUCCESS.labels(label).inc();
       return result;
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Delete a list of raw key-value pair from TiKV if key exists
-   *
-   * @param keys list of raw key
-   */
+  @Override
   public void batchDelete(List<ByteString> keys) {
     batchDelete(keys, false);
   }
 
-  /**
-   * Delete a list of raw key-value pair from TiKV if key exists, this API is atomic
-   *
-   * @param keys list of raw key
-   */
+  @Override
   public void batchDeleteAtomic(List<ByteString> keys) {
     batchDelete(keys, true);
   }
@@ -302,52 +311,75 @@ public class RawKVClient implements AutoCloseable {
   private void batchDelete(List<ByteString> keys, boolean atomic) {
     String label = "client_raw_batch_delete";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVBatchWriteSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "batchDelete");
+                put("keySize", String.valueOf(keys.size()));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVBatchWriteTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
-      doSendBatchDelete(backOffer, keys, atomic);
+      long deadline = System.currentTimeMillis() + conf.getRawKVBatchWriteTimeoutInMS();
+      doSendBatchDelete(backOffer, keys, atomic, deadline);
       RAW_REQUEST_SUCCESS.labels(label).inc();
       return;
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Get the TTL of a raw key from TiKV if key exists
-   *
-   * @param key raw key
-   * @return a Long indicating the TTL of key ttl is a non-null long value indicating TTL if key
-   *     exists. - ttl=0 if the key will never be outdated. - ttl=null if the key does not exist
-   */
+  @Override
   public Long getKeyTTL(ByteString key) {
     String label = "client_raw_get_key_ttl";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVReadSlowLogInMS(),
+            new HashMap<String, String>(2) {
+              {
+                put("func", "getKeyTTL");
+                put("key", KeyUtils.formatBytesUTF8(key));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVReadTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
       while (true) {
-        RegionStoreClient client = clientBuilder.build(key);
-        try {
-          Long result = client.rawGetKeyTTL(defaultBackOff(), key);
+        try (RegionStoreClient client = clientBuilder.build(key, backOffer)) {
+          slowLog.addProperty("region", client.getRegion().toString());
+          Long result = client.rawGetKeyTTL(backOffer, key);
           RAW_REQUEST_SUCCESS.labels(label).inc();
           return result;
         } catch (final TiKVException e) {
           backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+          logger.warn("Retry for getKeyTTL error", e);
         }
       }
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
+  @Override
   public List<List<KvPair>> batchScan(List<ScanOption> ranges) {
     String label = "client_raw_batch_scan";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    long deadline = System.currentTimeMillis() + conf.getRawKVScanTimeoutInMS();
+    List<Future<Pair<Integer, List<KvPair>>>> futureList = new ArrayList<>();
     try {
       if (ranges.isEmpty()) {
         return new ArrayList<>();
@@ -357,7 +389,7 @@ public class RawKVClient implements AutoCloseable {
       int num = 0;
       for (ScanOption scanOption : ranges) {
         int i = num;
-        completionService.submit(() -> Pair.create(i, scan(scanOption)));
+        futureList.add(completionService.submit(() -> Pair.create(i, scan(scanOption))));
         ++num;
       }
       List<List<KvPair>> scanResults = new ArrayList<>();
@@ -366,14 +398,16 @@ public class RawKVClient implements AutoCloseable {
       }
       for (int i = 0; i < num; i++) {
         try {
-          Pair<Integer, List<KvPair>> scanResult =
-              completionService.take().get(BackOffer.RAWKV_MAX_BACKOFF, TimeUnit.SECONDS);
+          Future<Pair<Integer, List<KvPair>>> future =
+              completionService.poll(deadline - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+          if (future == null) {
+            throw new TiKVException("TimeOut Exceeded for current operation.");
+          }
+          Pair<Integer, List<KvPair>> scanResult = future.get();
           scanResults.set(scanResult.first, scanResult.second);
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
           throw new TiKVException("Current thread interrupted.", e);
-        } catch (TimeoutException e) {
-          throw new TiKVException("TimeOut Exceeded for current operation. ", e);
         } catch (ExecutionException e) {
           throw new TiKVException("Execution exception met.", e);
         }
@@ -382,115 +416,115 @@ public class RawKVClient implements AutoCloseable {
       return scanResults;
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      for (Future<Pair<Integer, List<KvPair>>> future : futureList) {
+        future.cancel(true);
+      }
       throw e;
     } finally {
       requestTimer.observeDuration();
     }
   }
 
-  /**
-   * Scan raw key-value pairs from TiKV in range [startKey, endKey)
-   *
-   * @param startKey raw start key, inclusive
-   * @param endKey raw end key, exclusive
-   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, ByteString endKey, int limit) {
     return scan(startKey, endKey, limit, false);
   }
 
-  /**
-   * Scan raw key-value pairs from TiKV in range [startKey, endKey)
-   *
-   * @param startKey raw start key, inclusive
-   * @param endKey raw end key, exclusive
-   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
-   * @param keyOnly whether to scan in key-only mode
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, ByteString endKey, int limit, boolean keyOnly) {
     String label = "client_raw_scan";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVScanSlowLogInMS(),
+            new HashMap<String, String>(5) {
+              {
+                put("func", "scan");
+                put("startKey", KeyUtils.formatBytesUTF8(startKey));
+                put("endKey", KeyUtils.formatBytesUTF8(endKey));
+                put("limit", String.valueOf(limit));
+                put("keyOnly", String.valueOf(keyOnly));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVScanTimeoutInMS(), slowLog);
     try {
       Iterator<KvPair> iterator =
-          rawScanIterator(conf, clientBuilder, startKey, endKey, limit, keyOnly);
+          rawScanIterator(conf, clientBuilder, startKey, endKey, limit, keyOnly, backOffer);
       List<KvPair> result = new ArrayList<>();
       iterator.forEachRemaining(result::add);
       RAW_REQUEST_SUCCESS.labels(label).inc();
       return result;
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Scan raw key-value pairs from TiKV in range [startKey, ♾)
-   *
-   * @param startKey raw start key, inclusive
-   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, int limit) {
     return scan(startKey, limit, false);
   }
 
-  /**
-   * Scan raw key-value pairs from TiKV in range [startKey, ♾)
-   *
-   * @param startKey raw start key, inclusive
-   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
-   * @param keyOnly whether to scan in key-only mode
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, int limit, boolean keyOnly) {
     return scan(startKey, ByteString.EMPTY, limit, keyOnly);
   }
 
-  /**
-   * Scan all raw key-value pairs from TiKV in range [startKey, endKey)
-   *
-   * @param startKey raw start key, inclusive
-   * @param endKey raw end key, exclusive
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, ByteString endKey) {
     return scan(startKey, endKey, false);
   }
 
-  /**
-   * Scan all raw key-value pairs from TiKV in range [startKey, endKey)
-   *
-   * @param startKey raw start key, inclusive
-   * @param endKey raw end key, exclusive
-   * @param keyOnly whether to scan in key-only mode
-   * @return list of key-value pairs in range
-   */
+  @Override
   public List<KvPair> scan(ByteString startKey, ByteString endKey, boolean keyOnly) {
     String label = "client_raw_scan_without_limit";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVScanSlowLogInMS(),
+            new HashMap<String, String>(4) {
+              {
+                put("func", "scan");
+                put("startKey", KeyUtils.formatBytesUTF8(startKey));
+                put("endKey", KeyUtils.formatBytesUTF8(endKey));
+                put("keyOnly", String.valueOf(keyOnly));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVScanTimeoutInMS(), slowLog);
     try {
+      ByteString newStartKey = startKey;
       List<KvPair> result = new ArrayList<>();
       while (true) {
         Iterator<KvPair> iterator =
             rawScanIterator(
-                conf, clientBuilder, startKey, endKey, conf.getScanBatchSize(), keyOnly);
+                conf,
+                clientBuilder,
+                newStartKey,
+                endKey,
+                conf.getScanBatchSize(),
+                keyOnly,
+                backOffer);
         if (!iterator.hasNext()) {
           break;
         }
         iterator.forEachRemaining(result::add);
-        startKey = Key.toRawKey(result.get(result.size() - 1).getKey()).next().toByteString();
+        newStartKey = Key.toRawKey(result.get(result.size() - 1).getKey()).next().toByteString();
       }
       RAW_REQUEST_SUCCESS.labels(label).inc();
       return result;
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
@@ -502,69 +536,78 @@ public class RawKVClient implements AutoCloseable {
     return scan(startKey, endKey, limit, keyOnly);
   }
 
-  /**
-   * Scan keys with prefix
-   *
-   * @param prefixKey prefix key
-   * @param limit limit of keys retrieved
-   * @param keyOnly whether to scan in keyOnly mode
-   * @return kvPairs with the specified prefix
-   */
+  @Override
   public List<KvPair> scanPrefix(ByteString prefixKey, int limit, boolean keyOnly) {
     return scan(prefixKey, Key.toRawKey(prefixKey).nextPrefix().toByteString(), limit, keyOnly);
   }
 
+  @Override
   public List<KvPair> scanPrefix(ByteString prefixKey) {
     return scan(prefixKey, Key.toRawKey(prefixKey).nextPrefix().toByteString());
   }
 
+  @Override
   public List<KvPair> scanPrefix(ByteString prefixKey, boolean keyOnly) {
     return scan(prefixKey, Key.toRawKey(prefixKey).nextPrefix().toByteString(), keyOnly);
   }
 
-  /**
-   * Delete a raw key-value pair from TiKV if key exists
-   *
-   * @param key raw key to be deleted
-   */
+  @Override
   public void delete(ByteString key) {
+    delete(key, false);
+  }
+
+  @Override
+  public void deleteAtomic(ByteString key) {
+    delete(key, true);
+  }
+
+  private void delete(ByteString key, boolean atomic) {
     String label = "client_raw_delete";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    SlowLog slowLog =
+        new SlowLogImpl(
+            conf.getRawKVWriteSlowLogInMS(),
+            new HashMap<String, String>(3) {
+              {
+                put("func", "delete");
+                put("key", KeyUtils.formatBytesUTF8(key));
+                put("atomic", String.valueOf(atomic));
+              }
+            });
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(conf.getRawKVWriteTimeoutInMS(), slowLog);
     try {
-      BackOffer backOffer = defaultBackOff();
       while (true) {
-        RegionStoreClient client = clientBuilder.build(key);
-        try {
-          client.rawDelete(defaultBackOff(), key);
+        try (RegionStoreClient client = clientBuilder.build(key, backOffer)) {
+          slowLog.addProperty("region", client.getRegion().toString());
+          client.rawDelete(backOffer, key, atomic);
           RAW_REQUEST_SUCCESS.labels(label).inc();
           return;
         } catch (final TiKVException e) {
           backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e);
+          logger.warn("Retry for delete error", e);
         }
       }
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
+      slowLog.setError(e);
       throw e;
     } finally {
       requestTimer.observeDuration();
+      slowLog.log();
     }
   }
 
-  /**
-   * Delete all raw key-value pairs in range [startKey, endKey) from TiKV
-   *
-   * <p>Cautious, this API cannot be used concurrently, if multiple clients write keys into this
-   * range along with deleteRange API, the result will be undefined.
-   *
-   * @param startKey raw start key to be deleted
-   * @param endKey raw start key to be deleted
-   */
+  @Override
   public synchronized void deleteRange(ByteString startKey, ByteString endKey) {
     String label = "client_raw_delete_range";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
+    ConcreteBackOffer backOffer =
+        ConcreteBackOffer.newDeadlineBackOff(
+            conf.getRawKVCleanTimeoutInMS(), SlowLogEmptyImpl.INSTANCE);
     try {
-      BackOffer backOffer = defaultBackOff();
-      doSendDeleteRange(backOffer, startKey, endKey);
+      long deadline = System.currentTimeMillis() + conf.getRawKVCleanTimeoutInMS();
+      doSendDeleteRange(backOffer, startKey, endKey, deadline);
       RAW_REQUEST_SUCCESS.labels(label).inc();
     } catch (Exception e) {
       RAW_REQUEST_FAILURE.labels(label).inc();
@@ -574,24 +617,23 @@ public class RawKVClient implements AutoCloseable {
     }
   }
 
-  /**
-   * Delete all raw key-value pairs with the prefix `key` from TiKV
-   *
-   * <p>Cautious, this API cannot be used concurrently, if multiple clients write keys into this
-   * range along with deleteRange API, the result will be undefined.
-   *
-   * @param key prefix of keys to be deleted
-   */
+  @Override
   public synchronized void deletePrefix(ByteString key) {
     ByteString endKey = Key.toRawKey(key).nextPrefix().toByteString();
     deleteRange(key, endKey);
   }
 
   private void doSendBatchPut(
-      BackOffer backOffer, Map<ByteString, ByteString> kvPairs, long ttl, boolean atomic) {
+      BackOffer backOffer,
+      Map<ByteString, ByteString> kvPairs,
+      long ttl,
+      boolean atomic,
+      long deadline) {
     ExecutorCompletionService<List<Batch>> completionService =
         new ExecutorCompletionService<>(batchPutThreadPool);
 
+    List<Future<List<Batch>>> futureList = new ArrayList<>();
+
     Map<TiRegion, List<ByteString>> groupKeys =
         groupKeysByRegion(clientBuilder.getRegionManager(), kvPairs.keySet(), backOffer);
     List<Batch> batches = new ArrayList<>();
@@ -612,22 +654,32 @@ public class RawKVClient implements AutoCloseable {
     while (!taskQueue.isEmpty()) {
       List<Batch> task = taskQueue.poll();
       for (Batch batch : task) {
-        completionService.submit(
-            () -> doSendBatchPutInBatchesWithRetry(batch.getBackOffer(), batch, ttl, atomic));
+        futureList.add(
+            completionService.submit(
+                () -> doSendBatchPutInBatchesWithRetry(batch.getBackOffer(), batch, ttl, atomic)));
+      }
+
+      try {
+        getTasks(completionService, taskQueue, task, deadline - System.currentTimeMillis());
+      } catch (Exception e) {
+        for (Future<List<Batch>> future : futureList) {
+          future.cancel(true);
+        }
+        throw e;
       }
-      getTasks(completionService, taskQueue, task, BackOffer.RAWKV_MAX_BACKOFF);
     }
   }
 
   private List<Batch> doSendBatchPutInBatchesWithRetry(
       BackOffer backOffer, Batch batch, long ttl, boolean atomic) {
-    try (RegionStoreClient client = clientBuilder.build(batch.getRegion())) {
+    try (RegionStoreClient client = clientBuilder.build(batch.getRegion(), backOffer)) {
+      client.setTimeout(conf.getRawKVBatchWriteTimeoutInMS());
       client.rawBatchPut(backOffer, batch, ttl, atomic);
       return new ArrayList<>();
     } 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.warn("ReSplitting ranges for BatchPutRequest", e);
       // retry
       return doSendBatchPutWithRefetchRegion(backOffer, batch);
     }
@@ -652,10 +704,12 @@ public class RawKVClient implements AutoCloseable {
     return retryBatches;
   }
 
-  private List<KvPair> doSendBatchGet(BackOffer backOffer, List<ByteString> keys) {
+  private List<KvPair> doSendBatchGet(BackOffer backOffer, List<ByteString> keys, long deadline) {
     ExecutorCompletionService<Pair<List<Batch>, List<KvPair>>> completionService =
         new ExecutorCompletionService<>(batchGetThreadPool);
 
+    List<Future<Pair<List<Batch>, List<KvPair>>>> futureList = new ArrayList<>();
+
     List<Batch> batches =
         getBatches(backOffer, keys, RAW_BATCH_GET_SIZE, MAX_RAW_BATCH_LIMIT, this.clientBuilder);
 
@@ -666,11 +720,20 @@ public class RawKVClient implements AutoCloseable {
     while (!taskQueue.isEmpty()) {
       List<Batch> task = taskQueue.poll();
       for (Batch batch : task) {
-        completionService.submit(
-            () -> doSendBatchGetInBatchesWithRetry(batch.getBackOffer(), batch));
+        futureList.add(
+            completionService.submit(
+                () -> doSendBatchGetInBatchesWithRetry(batch.getBackOffer(), batch)));
+      }
+      try {
+        result.addAll(
+            getTasksWithOutput(
+                completionService, taskQueue, task, deadline - System.currentTimeMillis()));
+      } catch (Exception e) {
+        for (Future<Pair<List<Batch>, List<KvPair>>> future : futureList) {
+          future.cancel(true);
+        }
+        throw e;
       }
-      result.addAll(
-          getTasksWithOutput(completionService, taskQueue, task, BackOffer.RAWKV_MAX_BACKOFF));
     }
 
     return result;
@@ -678,8 +741,8 @@ public class RawKVClient implements AutoCloseable {
 
   private Pair<List<Batch>, List<KvPair>> doSendBatchGetInBatchesWithRetry(
       BackOffer backOffer, Batch batch) {
-    RegionStoreClient client = clientBuilder.build(batch.getRegion());
-    try {
+
+    try (RegionStoreClient client = clientBuilder.build(batch.getRegion(), backOffer)) {
       List<KvPair> partialResult = client.rawBatchGet(backOffer, batch.getKeys());
       return Pair.create(new ArrayList<>(), partialResult);
     } catch (final TiKVException e) {
@@ -697,10 +760,13 @@ public class RawKVClient implements AutoCloseable {
         backOffer, batch.getKeys(), RAW_BATCH_GET_SIZE, MAX_RAW_BATCH_LIMIT, clientBuilder);
   }
 
-  private void doSendBatchDelete(BackOffer backOffer, List<ByteString> keys, boolean atomic) {
+  private void doSendBatchDelete(
+      BackOffer backOffer, List<ByteString> keys, boolean atomic, long deadline) {
     ExecutorCompletionService<List<Batch>> completionService =
         new ExecutorCompletionService<>(batchDeleteThreadPool);
 
+    List<Future<List<Batch>>> futureList = new ArrayList<>();
+
     List<Batch> batches =
         getBatches(backOffer, keys, RAW_BATCH_DELETE_SIZE, MAX_RAW_BATCH_LIMIT, this.clientBuilder);
 
@@ -710,17 +776,24 @@ public class RawKVClient implements AutoCloseable {
     while (!taskQueue.isEmpty()) {
       List<Batch> task = taskQueue.poll();
       for (Batch batch : task) {
-        completionService.submit(
-            () -> doSendBatchDeleteInBatchesWithRetry(batch.getBackOffer(), batch, atomic));
+        futureList.add(
+            completionService.submit(
+                () -> doSendBatchDeleteInBatchesWithRetry(batch.getBackOffer(), batch, atomic)));
+      }
+      try {
+        getTasks(completionService, taskQueue, task, deadline - System.currentTimeMillis());
+      } catch (Exception e) {
+        for (Future<List<Batch>> future : futureList) {
+          future.cancel(true);
+        }
+        throw e;
       }
-      getTasks(completionService, taskQueue, task, BackOffer.RAWKV_MAX_BACKOFF);
     }
   }
 
   private List<Batch> doSendBatchDeleteInBatchesWithRetry(
       BackOffer backOffer, Batch batch, boolean atomic) {
-    RegionStoreClient client = clientBuilder.build(batch.getRegion());
-    try {
+    try (RegionStoreClient client = clientBuilder.build(batch.getRegion(), backOffer)) {
       client.rawBatchDelete(backOffer, batch.getKeys(), atomic);
       return new ArrayList<>();
     } catch (final TiKVException e) {
@@ -745,10 +818,13 @@ public class RawKVClient implements AutoCloseable {
     return key2;
   }
 
-  private void doSendDeleteRange(BackOffer backOffer, ByteString startKey, ByteString endKey) {
+  private void doSendDeleteRange(
+      BackOffer backOffer, ByteString startKey, ByteString endKey, long deadline) {
     ExecutorCompletionService<List<DeleteRange>> completionService =
         new ExecutorCompletionService<>(deleteRangeThreadPool);
 
+    List<Future<List<DeleteRange>>> futureList = new ArrayList<>();
+
     List<TiRegion> regions = fetchRegionsFromRange(backOffer, startKey, endKey);
     List<DeleteRange> ranges = new ArrayList<>();
     for (int i = 0; i < regions.size(); i++) {
@@ -762,14 +838,23 @@ public class RawKVClient implements AutoCloseable {
     while (!taskQueue.isEmpty()) {
       List<DeleteRange> task = taskQueue.poll();
       for (DeleteRange range : task) {
-        completionService.submit(() -> doSendDeleteRangeWithRetry(range.getBackOffer(), range));
+        futureList.add(
+            completionService.submit(
+                () -> doSendDeleteRangeWithRetry(range.getBackOffer(), range)));
+      }
+      try {
+        getTasks(completionService, taskQueue, task, deadline - System.currentTimeMillis());
+      } catch (Exception e) {
+        for (Future<List<DeleteRange>> future : futureList) {
+          future.cancel(true);
+        }
+        throw e;
       }
-      getTasks(completionService, taskQueue, task, BackOffer.RAWKV_MAX_BACKOFF);
     }
   }
 
   private List<DeleteRange> doSendDeleteRangeWithRetry(BackOffer backOffer, DeleteRange range) {
-    try (RegionStoreClient client = clientBuilder.build(range.getRegion())) {
+    try (RegionStoreClient client = clientBuilder.build(range.getRegion(), backOffer)) {
       client.setTimeout(conf.getScanTimeout());
       client.rawDeleteRange(backOffer, range.getStartKey(), range.getEndKey());
       return new ArrayList<>();
@@ -810,7 +895,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();
@@ -827,14 +914,11 @@ public class RawKVClient implements AutoCloseable {
       ByteString startKey,
       ByteString endKey,
       int limit,
-      boolean keyOnly) {
+      boolean keyOnly,
+      BackOffer backOffer) {
     if (limit > MAX_RAW_SCAN_LIMIT) {
       throw ERR_MAX_SCAN_LIMIT_EXCEEDED;
     }
-    return new RawScanIterator(conf, builder, startKey, endKey, limit, keyOnly);
-  }
-
-  private BackOffer defaultBackOff() {
-    return ConcreteBackOffer.newRawKVBackOff();
+    return new RawScanIterator(conf, builder, startKey, endKey, limit, keyOnly, backOffer);
   }
 }

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

@@ -0,0 +1,257 @@
+/*
+ * 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.raw;
+
+import com.google.protobuf.ByteString;
+import java.util.List;
+import java.util.Map;
+import org.tikv.common.util.ScanOption;
+import org.tikv.kvproto.Kvrpcpb;
+
+public interface RawKVClientBase extends AutoCloseable {
+  // https://www.github.com/pingcap/tidb/blob/master/store/tikv/rawkv.go
+  int MAX_RAW_SCAN_LIMIT = 10240;
+  int MAX_RAW_BATCH_LIMIT = 1024;
+  int RAW_BATCH_PUT_SIZE = 1024 * 1024;
+  int RAW_BATCH_GET_SIZE = 16 * 1024;
+  int RAW_BATCH_DELETE_SIZE = 16 * 1024;
+
+  /**
+   * Put a raw key-value pair to TiKV
+   *
+   * @param key raw key
+   * @param value raw value
+   */
+  void put(ByteString key, ByteString value);
+
+  /**
+   * Put a raw key-value pair to TiKV
+   *
+   * @param key raw key
+   * @param value raw value
+   * @param ttl the ttl of the key (in seconds), 0 means the key will never be outdated
+   */
+  void put(ByteString key, ByteString value, long ttl);
+
+  /**
+   * Put a raw key-value pair to TiKV. This API is atomic.
+   *
+   * @param key raw key
+   * @param value raw value
+   * @param ttl the ttl of the key (in seconds), 0 means the key will never be outdated
+   */
+  void putAtomic(ByteString key, ByteString value, long ttl);
+
+  /**
+   * Put a key-value pair if it does not exist. This API is atomic.
+   *
+   * @param key key
+   * @param value value
+   * @return a ByteString. returns ByteString.EMPTY if the value is written successfully. returns
+   *     the previous key if the value already exists, and does not write to TiKV.
+   */
+  ByteString putIfAbsent(ByteString key, ByteString value);
+
+  /**
+   * Put a key-value pair with TTL if it does not exist. This API is atomic.
+   *
+   * @param key key
+   * @param value value
+   * @param ttl TTL of key (in seconds), 0 means the key will never be outdated.
+   * @return a ByteString. returns ByteString.EMPTY if the value is written successfully. returns
+   *     the previous key if the value already exists, and does not write to TiKV.
+   */
+  ByteString putIfAbsent(ByteString key, ByteString value, long ttl);
+
+  /**
+   * Put a set of raw key-value pair to TiKV, this API does not ensure the operation is atomic.
+   *
+   * @param kvPairs kvPairs
+   */
+  void batchPut(Map<ByteString, ByteString> kvPairs);
+
+  /**
+   * Put a set of raw key-value pair to TiKV, this API does not ensure the operation is atomic.
+   *
+   * @param kvPairs kvPairs
+   * @param ttl the TTL of keys to be put (in seconds), 0 means the keys will never be outdated
+   */
+  void batchPut(Map<ByteString, ByteString> kvPairs, long ttl);
+
+  /**
+   * Put a set of raw key-value pair to TiKV, this API is atomic
+   *
+   * @param kvPairs kvPairs
+   */
+  void batchPutAtomic(Map<ByteString, ByteString> kvPairs);
+
+  /**
+   * Put a set of raw key-value pair to TiKV, this API is atomic.
+   *
+   * @param kvPairs kvPairs
+   * @param ttl the TTL of keys to be put (in seconds), 0 means the keys will never be outdated
+   */
+  void batchPutAtomic(Map<ByteString, ByteString> kvPairs, long ttl);
+
+  /**
+   * Get a raw key-value pair from TiKV if key exists
+   *
+   * @param key raw key
+   * @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
+   */
+  ByteString get(ByteString key);
+
+  /**
+   * Get a list of raw key-value pair from TiKV if key exists
+   *
+   * @param keys list of raw key
+   * @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
+   */
+  List<Kvrpcpb.KvPair> batchGet(List<ByteString> keys);
+
+  /**
+   * Delete a list of raw key-value pair from TiKV if key exists
+   *
+   * @param keys list of raw key
+   */
+  void batchDelete(List<ByteString> keys);
+
+  /**
+   * Delete a list of raw key-value pair from TiKV if key exists, this API is atomic
+   *
+   * @param keys list of raw key
+   */
+  void batchDeleteAtomic(List<ByteString> keys);
+
+  /**
+   * Get the TTL of a raw key from TiKV if key exists
+   *
+   * @param key raw key
+   * @return a Long indicating the TTL of key ttl is a non-null long value indicating TTL if key
+   *     exists. - ttl=0 if the key will never be outdated. - ttl=null if the key does not exist
+   */
+  Long getKeyTTL(ByteString key);
+
+  List<List<Kvrpcpb.KvPair>> batchScan(List<ScanOption> ranges);
+
+  /**
+   * Scan raw key-value pairs from TiKV in range [startKey, endKey)
+   *
+   * @param startKey raw start key, inclusive
+   * @param endKey raw end key, exclusive
+   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, int limit);
+
+  /**
+   * Scan raw key-value pairs from TiKV in range [startKey, endKey)
+   *
+   * @param startKey raw start key, inclusive
+   * @param endKey raw end key, exclusive
+   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
+   * @param keyOnly whether to scan in key-only mode
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, int limit, boolean keyOnly);
+
+  /**
+   * Scan raw key-value pairs from TiKV in range [startKey, ♾)
+   *
+   * @param startKey raw start key, inclusive
+   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit);
+
+  /**
+   * Scan raw key-value pairs from TiKV in range [startKey, ♾)
+   *
+   * @param startKey raw start key, inclusive
+   * @param limit limit of key-value pairs scanned, should be less than {@link #MAX_RAW_SCAN_LIMIT}
+   * @param keyOnly whether to scan in key-only mode
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit, boolean keyOnly);
+
+  /**
+   * Scan all raw key-value pairs from TiKV in range [startKey, endKey)
+   *
+   * @param startKey raw start key, inclusive
+   * @param endKey raw end key, exclusive
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey);
+
+  /**
+   * Scan all raw key-value pairs from TiKV in range [startKey, endKey)
+   *
+   * @param startKey raw start key, inclusive
+   * @param endKey raw end key, exclusive
+   * @param keyOnly whether to scan in key-only mode
+   * @return list of key-value pairs in range
+   */
+  List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, boolean keyOnly);
+
+  /**
+   * Scan keys with prefix
+   *
+   * @param prefixKey prefix key
+   * @param limit limit of keys retrieved
+   * @param keyOnly whether to scan in keyOnly mode
+   * @return kvPairs with the specified prefix
+   */
+  List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, int limit, boolean keyOnly);
+
+  List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey);
+
+  List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, boolean keyOnly);
+
+  /**
+   * Delete a raw key-value pair from TiKV if key exists
+   *
+   * @param key raw key to be deleted
+   */
+  void delete(ByteString key);
+
+  /**
+   * Delete a raw key-value pair from TiKV if key exists. This API is atomic.
+   *
+   * @param key raw key to be deleted
+   */
+  void deleteAtomic(ByteString key);
+
+  /**
+   * Delete all raw key-value pairs in range [startKey, endKey) from TiKV
+   *
+   * <p>Cautious, this API cannot be used concurrently, if multiple clients write keys into this
+   * range along with deleteRange API, the result will be undefined.
+   *
+   * @param startKey raw start key to be deleted
+   * @param endKey raw start key to be deleted
+   */
+  void deleteRange(ByteString startKey, ByteString endKey);
+
+  /**
+   * Delete all raw key-value pairs with the prefix `key` from TiKV
+   *
+   * <p>Cautious, this API cannot be used concurrently, if multiple clients write keys into this
+   * range along with deleteRange API, the result will be undefined.
+   *
+   * @param key prefix of keys to be deleted
+   */
+  void deletePrefix(ByteString key);
+}

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

@@ -0,0 +1,277 @@
+/*
+ * 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.raw;
+
+import com.google.protobuf.ByteString;
+import io.prometheus.client.Counter;
+import io.prometheus.client.Histogram;
+import java.util.List;
+import java.util.Map;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.exception.CircuitBreakerOpenException;
+import org.tikv.common.util.ScanOption;
+import org.tikv.kvproto.Kvrpcpb;
+import org.tikv.service.failsafe.CircuitBreaker;
+
+public class SmartRawKVClient implements RawKVClientBase {
+  private static final Logger logger = LoggerFactory.getLogger(SmartRawKVClient.class);
+
+  private static final Histogram REQUEST_LATENCY =
+      Histogram.build()
+          .name("client_java_smart_raw_requests_latency")
+          .help("client smart raw request latency.")
+          .labelNames("type")
+          .register();
+
+  private static final Counter REQUEST_SUCCESS =
+      Counter.build()
+          .name("client_java_smart_raw_requests_success")
+          .help("client smart raw request success.")
+          .labelNames("type")
+          .register();
+
+  private static final Counter REQUEST_FAILURE =
+      Counter.build()
+          .name("client_java_smart_raw_requests_failure")
+          .help("client smart raw request failure.")
+          .labelNames("type")
+          .register();
+
+  private static final Counter CIRCUIT_BREAKER_OPENED =
+      Counter.build()
+          .name("client_java_smart_raw_circuit_breaker_opened")
+          .help("client smart raw circuit breaker opened.")
+          .labelNames("type")
+          .register();
+
+  private final RawKVClientBase client;
+  private final CircuitBreaker circuitBreaker;
+
+  public SmartRawKVClient(RawKVClientBase client, CircuitBreaker breaker) {
+    this.client = client;
+    this.circuitBreaker = breaker;
+  }
+
+  @Override
+  public void put(ByteString key, ByteString value) {
+    callWithCircuitBreaker("put", () -> client.put(key, value));
+  }
+
+  @Override
+  public void put(ByteString key, ByteString value, long ttl) {
+    callWithCircuitBreaker("put", () -> client.put(key, value, ttl));
+  }
+
+  @Override
+  public void putAtomic(ByteString key, ByteString value, long ttl) {
+    callWithCircuitBreaker("putAtomic", () -> client.putAtomic(key, value, ttl));
+  }
+
+  @Override
+  public ByteString putIfAbsent(ByteString key, ByteString value) {
+    return callWithCircuitBreaker("putIfAbsent", () -> client.putIfAbsent(key, value));
+  }
+
+  @Override
+  public ByteString putIfAbsent(ByteString key, ByteString value, long ttl) {
+    return callWithCircuitBreaker("putIfAbsent", () -> client.putIfAbsent(key, value, ttl));
+  }
+
+  @Override
+  public void batchPut(Map<ByteString, ByteString> kvPairs) {
+    callWithCircuitBreaker("batchPut", () -> client.batchPut(kvPairs));
+  }
+
+  @Override
+  public void batchPut(Map<ByteString, ByteString> kvPairs, long ttl) {
+    callWithCircuitBreaker("batchPut", () -> client.batchPut(kvPairs, ttl));
+  }
+
+  @Override
+  public void batchPutAtomic(Map<ByteString, ByteString> kvPairs) {
+    callWithCircuitBreaker("batchPutAtomic", () -> client.batchPutAtomic(kvPairs));
+  }
+
+  @Override
+  public void batchPutAtomic(Map<ByteString, ByteString> kvPairs, long ttl) {
+    callWithCircuitBreaker("batchPutAtomic", () -> client.batchPutAtomic(kvPairs, ttl));
+  }
+
+  @Override
+  public ByteString get(ByteString key) {
+    return callWithCircuitBreaker("get", () -> client.get(key));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> batchGet(List<ByteString> keys) {
+    return callWithCircuitBreaker("batchGet", () -> client.batchGet(keys));
+  }
+
+  @Override
+  public void batchDelete(List<ByteString> keys) {
+    callWithCircuitBreaker("batchDelete", () -> client.batchDelete(keys));
+  }
+
+  @Override
+  public void batchDeleteAtomic(List<ByteString> keys) {
+    callWithCircuitBreaker("batchDeleteAtomic", () -> client.batchDeleteAtomic(keys));
+  }
+
+  @Override
+  public Long getKeyTTL(ByteString key) {
+    return callWithCircuitBreaker("getKeyTTL", () -> client.getKeyTTL(key));
+  }
+
+  @Override
+  public List<List<Kvrpcpb.KvPair>> batchScan(List<ScanOption> ranges) {
+    return callWithCircuitBreaker("batchScan", () -> client.batchScan(ranges));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, int limit) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, limit));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(
+      ByteString startKey, ByteString endKey, int limit, boolean keyOnly) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, limit, keyOnly));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, limit));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit, boolean keyOnly) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, limit, keyOnly));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, boolean keyOnly) {
+    return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, keyOnly));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, int limit, boolean keyOnly) {
+    return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey, limit, keyOnly));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey) {
+    return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey));
+  }
+
+  @Override
+  public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, boolean keyOnly) {
+    return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey, keyOnly));
+  }
+
+  @Override
+  public void delete(ByteString key) {
+    callWithCircuitBreaker("delete", () -> client.delete(key));
+  }
+
+  @Override
+  public void deleteAtomic(ByteString key) {
+    callWithCircuitBreaker("deleteAtomic", () -> client.deleteAtomic(key));
+  }
+
+  @Override
+  public void deleteRange(ByteString startKey, ByteString endKey) {
+    callWithCircuitBreaker("deleteRange", () -> client.deleteRange(startKey, endKey));
+  }
+
+  @Override
+  public void deletePrefix(ByteString key) {
+    callWithCircuitBreaker("deletePrefix", () -> client.deletePrefix(key));
+  }
+
+  <T> T callWithCircuitBreaker(String funcName, Function1<T> func) {
+    Histogram.Timer requestTimer = REQUEST_LATENCY.labels(funcName).startTimer();
+    try {
+      T result = callWithCircuitBreaker0(funcName, func);
+      REQUEST_SUCCESS.labels(funcName).inc();
+      return result;
+    } catch (Exception e) {
+      REQUEST_FAILURE.labels(funcName).inc();
+      throw e;
+    } finally {
+      requestTimer.observeDuration();
+    }
+  }
+
+  private <T> T callWithCircuitBreaker0(String funcName, Function1<T> func) {
+    if (circuitBreaker.allowRequest()) {
+      try {
+        T result = func.apply();
+        circuitBreaker.getMetrics().recordSuccess();
+        return result;
+      } catch (Exception e) {
+        circuitBreaker.getMetrics().recordFailure();
+        throw e;
+      }
+    } else if (circuitBreaker.attemptExecution()) {
+      logger.debug("attemptExecution");
+      try {
+        T result = func.apply();
+        circuitBreaker.getMetrics().recordSuccess();
+        circuitBreaker.recordAttemptSuccess();
+        logger.debug("markSuccess");
+        return result;
+      } catch (Exception e) {
+        circuitBreaker.getMetrics().recordFailure();
+        circuitBreaker.recordAttemptFailure();
+        logger.debug("markNonSuccess");
+        throw e;
+      }
+    } else {
+      logger.debug("Circuit Breaker Opened");
+      CIRCUIT_BREAKER_OPENED.labels(funcName).inc();
+      throw new CircuitBreakerOpenException();
+    }
+  }
+
+  private void callWithCircuitBreaker(String funcName, Function0 func) {
+    callWithCircuitBreaker(
+        funcName,
+        (Function1<Void>)
+            () -> {
+              func.apply();
+              return null;
+            });
+  }
+
+  @Override
+  public void close() throws Exception {
+    client.close();
+  }
+
+  public interface Function1<T> {
+    T apply();
+  }
+
+  public interface Function0 {
+    void apply();
+  }
+}

src/main/java/org/tikv/service/failsafe/CircuitBreaker.java

@@ -0,0 +1,60 @@
+/*
+ * 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.service.failsafe;
+
+import java.io.Closeable;
+
+public interface CircuitBreaker extends Closeable {
+
+  enum Status {
+    CLOSED(0),
+    HALF_OPEN(1),
+    OPEN(2);
+
+    private final int value;
+
+    private Status(int value) {
+      this.value = value;
+    }
+
+    public int getValue() {
+      return value;
+    }
+  }
+
+  /**
+   * Every requests asks this if it is allowed to proceed or not. It is idempotent and does not
+   * modify any internal state.
+   *
+   * @return boolean whether a request should be permitted
+   */
+  boolean allowRequest();
+
+  /**
+   * Invoked at start of command execution to attempt an execution. This is non-idempotent - it may
+   * modify internal state.
+   */
+  boolean attemptExecution();
+
+  /** Invoked on successful executions as part of feedback mechanism when in a half-open state. */
+  void recordAttemptSuccess();
+
+  /** Invoked on unsuccessful executions as part of feedback mechanism when in a half-open state. */
+  void recordAttemptFailure();
+
+  /** Get the Circuit Breaker Metrics Object. */
+  CircuitBreakerMetrics getMetrics();
+}

src/main/java/org/tikv/service/failsafe/CircuitBreakerImpl.java

@@ -0,0 +1,202 @@
+/*
+ * 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.service.failsafe;
+
+import io.prometheus.client.Counter;
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.tikv.common.TiConfiguration;
+
+public class CircuitBreakerImpl implements CircuitBreaker {
+  private static final Logger logger = LoggerFactory.getLogger(CircuitBreakerImpl.class);
+
+  private static final Counter CIRCUIT_BREAKER_ATTEMPT_COUNTER =
+      Counter.build()
+          .name("client_java_circuit_breaker_attempt_counter")
+          .help("client circuit breaker attempt counter.")
+          .labelNames("type")
+          .register();
+
+  private final boolean enable;
+  private final int windowInSeconds;
+  private final int errorThresholdPercentage;
+  private final int requestVolumeThreshold;
+  private final int sleepWindowInSeconds;
+  private final int attemptRequestCount;
+
+  private final AtomicLong circuitOpened = new AtomicLong(-1);
+  private final AtomicReference<Status> status = new AtomicReference<>(Status.CLOSED);
+  private final AtomicLong attemptCount = new AtomicLong(0);
+  private final AtomicLong attemptSuccessCount = new AtomicLong(0);
+
+  private final CircuitBreakerMetrics metrics;
+
+  public CircuitBreakerImpl(TiConfiguration conf) {
+    this(
+        conf.isCircuitBreakEnable(),
+        conf.getCircuitBreakAvailabilityWindowInSeconds(),
+        conf.getCircuitBreakAvailabilityErrorThresholdPercentage(),
+        conf.getCircuitBreakAvailabilityRequestVolumnThreshold(),
+        conf.getCircuitBreakSleepWindowInSeconds(),
+        conf.getCircuitBreakAttemptRequestCount());
+  }
+
+  public CircuitBreakerImpl(
+      boolean enable,
+      int windowInSeconds,
+      int errorThresholdPercentage,
+      int requestVolumeThreshold,
+      int sleepWindowInSeconds,
+      int attemptRequestCount) {
+    this.enable = enable;
+    this.windowInSeconds = windowInSeconds;
+    this.errorThresholdPercentage = errorThresholdPercentage;
+    this.requestVolumeThreshold = requestVolumeThreshold;
+    this.sleepWindowInSeconds = sleepWindowInSeconds;
+    this.attemptRequestCount = attemptRequestCount;
+    this.metrics =
+        enable ? new CircuitBreakerMetricsImpl(windowInSeconds) : new NoOpCircuitBreakerMetrics();
+    this.metrics.addListener(getMetricsListener());
+  }
+
+  private MetricsListener getMetricsListener() {
+    return hc -> {
+      logger.debug("onNext " + hc.toString());
+      // check if we are past the requestVolumeThreshold
+      if (hc.getTotalRequests() < requestVolumeThreshold) {
+        // we are not past the minimum volume threshold for the stat window,
+        // so no change to circuit status.
+        // if it was CLOSED, it stays CLOSED
+        // if it was half-open, we need to wait for some successful command executions
+        // if it was open, we need to wait for sleep window to elapse
+      } else {
+        if (hc.getErrorPercentage() < errorThresholdPercentage) {
+          // we are not past the minimum error threshold for the stat window,
+          // so no change to circuit status.
+          // if it was CLOSED, it stays CLOSED
+          // if it was half-open, we need to wait for some successful command executions
+          // if it was open, we need to wait for sleep window to elapse
+        } else {
+          // our failure rate is too high, we need to set the state to OPEN
+          close2Open();
+        }
+      }
+    };
+  }
+
+  @Override
+  public CircuitBreakerMetrics getMetrics() {
+    return metrics;
+  }
+
+  @Override
+  public boolean allowRequest() {
+    if (!enable) {
+      return true;
+    }
+    return !isOpen();
+  }
+
+  boolean isOpen() {
+    return circuitOpened.get() >= 0;
+  }
+
+  Status getStatus() {
+    return status.get();
+  }
+
+  @Override
+  public void recordAttemptSuccess() {
+    CIRCUIT_BREAKER_ATTEMPT_COUNTER.labels("success").inc();
+    if (attemptSuccessCount.incrementAndGet() >= this.attemptRequestCount) {
+      halfOpen2Close();
+    }
+  }
+
+  @Override
+  public void recordAttemptFailure() {
+    CIRCUIT_BREAKER_ATTEMPT_COUNTER.labels("failure").inc();
+    halfOpen2Open();
+  }
+
+  @Override
+  public boolean attemptExecution() {
+    if (allowRequest()) {
+      return true;
+    } else {
+      if (isAfterSleepWindow()) {
+        // only the `attemptRequestCount` requests after sleep window should execute
+        // if all the executing commands succeed, the status will transition to CLOSED
+        // if some of the executing commands fail, the status will transition to OPEN
+        open2HalfOpen();
+        return attemptCount.incrementAndGet() <= attemptRequestCount;
+      } else {
+        return false;
+      }
+    }
+  }
+
+  private boolean isAfterSleepWindow() {
+    final long circuitOpenTime = circuitOpened.get();
+    final long currentTime = System.currentTimeMillis();
+    final long sleepWindowTime = (long) sleepWindowInSeconds * 1000;
+    return currentTime >= circuitOpenTime + sleepWindowTime;
+  }
+
+  private void close2Open() {
+    if (status.compareAndSet(Status.CLOSED, Status.OPEN)) {
+      // This thread wins the race to open the circuit
+      // it sets the start time for the sleep window
+      circuitOpened.set(System.currentTimeMillis());
+      logger.info("CLOSED => OPEN");
+    }
+  }
+
+  private void halfOpen2Close() {
+    if (status.compareAndSet(Status.HALF_OPEN, Status.CLOSED)) {
+      // This thread wins the race to close the circuit
+      circuitOpened.set(-1L);
+      logger.info("HALF_OPEN => CLOSED");
+    }
+  }
+
+  private void open2HalfOpen() {
+    if (status.compareAndSet(Status.OPEN, Status.HALF_OPEN)) {
+      // This thread wins the race to half close the circuit
+      // it resets the attempt count
+      attemptCount.set(0);
+      attemptSuccessCount.set(0);
+      logger.info("OPEN => HALF_OPEN");
+    }
+  }
+
+  private void halfOpen2Open() {
+    if (status.compareAndSet(Status.HALF_OPEN, Status.OPEN)) {
+      // This thread wins the race to re-open the circuit
+      // it resets the start time for the sleep window
+      circuitOpened.set(System.currentTimeMillis());
+      logger.info("HALF_OPEN => OPEN");
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    metrics.close();
+  }
+}

src/main/java/org/tikv/service/failsafe/CircuitBreakerMetrics.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.service.failsafe;
+
+import java.io.Closeable;
+
+public interface CircuitBreakerMetrics extends Closeable {
+  /** Record a successful call. */
+  void recordSuccess();
+
+  /** Record a failure call. */
+  void recordFailure();
+
+  /** Add metrics listener. */
+  void addListener(MetricsListener metricsListener);
+}

src/main/java/org/tikv/service/failsafe/CircuitBreakerMetricsImpl.java

@@ -0,0 +1,120 @@
+/*
+ * 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.service.failsafe;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.commons.lang3.concurrent.BasicThreadFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CircuitBreakerMetricsImpl implements CircuitBreakerMetrics {
+  private static final Logger logger = LoggerFactory.getLogger(CircuitBreakerMetricsImpl.class);
+
+  private final int windowInMS;
+  private final List<MetricsListener> listeners;
+  private final AtomicReference<SingleWindowMetrics> currentMetrics;
+
+  private final ScheduledExecutorService scheduler;
+  private static final int SCHEDULER_INITIAL_DELAY = 1000;
+  private static final int SCHEDULER_PERIOD = 1000;
+
+  public CircuitBreakerMetricsImpl(int windowInSeconds) {
+    this.windowInMS = windowInSeconds * 1000;
+    this.listeners = new ArrayList<>();
+    this.currentMetrics = new AtomicReference<>(new SingleWindowMetrics());
+
+    scheduler =
+        new ScheduledThreadPoolExecutor(
+            1,
+            new BasicThreadFactory.Builder()
+                .namingPattern("circuit-breaker-metrics-%d")
+                .daemon(true)
+                .build());
+
+    scheduler.scheduleAtFixedRate(
+        this::onReachCircuitWindow,
+        SCHEDULER_INITIAL_DELAY,
+        SCHEDULER_PERIOD,
+        TimeUnit.MILLISECONDS);
+  }
+
+  @Override
+  public void recordSuccess() {
+    currentMetrics.get().recordSuccess();
+  }
+
+  @Override
+  public void recordFailure() {
+    currentMetrics.get().recordFailure();
+  }
+
+  private void onReachCircuitWindow() {
+    SingleWindowMetrics singleWindowMetrics = currentMetrics.get();
+    if (System.currentTimeMillis() < singleWindowMetrics.getStartMS() + windowInMS) {
+      return;
+    }
+    if (!currentMetrics.compareAndSet(singleWindowMetrics, new SingleWindowMetrics())) {
+      return;
+    }
+    logger.debug("window timeout, reset SingleWindowMetrics");
+    HealthCounts healthCounts = singleWindowMetrics.getHealthCounts();
+    for (MetricsListener metricsListener : listeners) {
+      metricsListener.onNext(healthCounts);
+    }
+  }
+
+  @Override
+  public void addListener(MetricsListener metricsListener) {
+    listeners.add(metricsListener);
+  }
+
+  @Override
+  public void close() throws IOException {
+    scheduler.shutdown();
+  }
+
+  /** Instead of using SingleWindowMetrics, it is better to use RollingWindowMetrics. */
+  static class SingleWindowMetrics {
+    private final long startMS = System.currentTimeMillis();
+    private final AtomicLong totalCount = new AtomicLong(0);
+    private final AtomicLong errorCount = new AtomicLong(0);
+
+    public void recordSuccess() {
+      totalCount.incrementAndGet();
+    }
+
+    public void recordFailure() {
+      totalCount.incrementAndGet();
+
+      errorCount.incrementAndGet();
+    }
+
+    public HealthCounts getHealthCounts() {
+      return new HealthCounts(totalCount.get(), errorCount.get());
+    }
+
+    public long getStartMS() {
+      return startMS;
+    }
+  }
+}

src/main/java/org/tikv/service/failsafe/HealthCounts.java

@@ -0,0 +1,56 @@
+/*
+ * 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.service.failsafe;
+
+public class HealthCounts {
+  private final long totalCount;
+  private final long errorCount;
+  private final int errorPercentage;
+
+  HealthCounts(long total, long error) {
+    this.totalCount = total;
+    this.errorCount = error;
+    if (totalCount > 0) {
+      this.errorPercentage = (int) ((double) errorCount / totalCount * 100);
+    } else {
+      this.errorPercentage = 0;
+    }
+  }
+
+  public long getTotalRequests() {
+    return totalCount;
+  }
+
+  public long getErrorCount() {
+    return errorCount;
+  }
+
+  public int getErrorPercentage() {
+    return errorPercentage;
+  }
+
+  @Override
+  public String toString() {
+    return "HealthCounts{"
+        + "totalCount="
+        + totalCount
+        + ", errorCount="
+        + errorCount
+        + ", errorPercentage="
+        + errorPercentage
+        + '}';
+  }
+}

src/main/java/org/tikv/service/failsafe/MetricsListener.java

@@ -0,0 +1,20 @@
+/*
+ * 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.service.failsafe;
+
+public interface MetricsListener {
+  void onNext(HealthCounts healthCounts);
+}

src/main/java/org/tikv/service/failsafe/NoOpCircuitBreakerMetrics.java

@@ -0,0 +1,40 @@
+/*
+ * 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.service.failsafe;
+
+import java.io.IOException;
+
+public class NoOpCircuitBreakerMetrics implements CircuitBreakerMetrics {
+  @Override
+  public void recordSuccess() {
+    // do nothing
+  }
+
+  @Override
+  public void recordFailure() {
+    // do nothing
+  }
+
+  @Override
+  public void addListener(MetricsListener metricsListener) {
+    // do nothing
+  }
+
+  @Override
+  public void close() throws IOException {
+    // do nothing
+  }
+}

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,
+      TikvGrpc.TikvFutureStub 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);
-    } else if (StoreVersion.compareTo(storeVersion, Version.RESOLVE_LOCK_V4) < 0) {
+          conf, region, store, blockingStub, asyncStub, channelFactory, regionManager);
+    } 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;
@@ -51,7 +52,7 @@ import org.tikv.kvproto.Kvrpcpb.ResolveLockRequest;
 import org.tikv.kvproto.Kvrpcpb.ResolveLockResponse;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
-import org.tikv.kvproto.TikvGrpc.TikvStub;
+import org.tikv.kvproto.TikvGrpc.TikvFutureStub;
 
 /** Before v3.0.5 TiDB uses the ttl on secondary lock. */
 public class LockResolverClientV2 extends AbstractRegionStoreClient
@@ -74,11 +75,12 @@ public class LockResolverClientV2 extends AbstractRegionStoreClient
   public LockResolverClientV2(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
-      TikvStub asyncStub,
+      TikvFutureStub 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.RegionManager;
-import org.tikv.common.region.RegionStoreClient;
-import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.*;
 import org.tikv.common.region.TiRegion.RegionVerID;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
@@ -52,7 +49,7 @@ import org.tikv.kvproto.Kvrpcpb.CleanupRequest;
 import org.tikv.kvproto.Kvrpcpb.CleanupResponse;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
-import org.tikv.kvproto.TikvGrpc.TikvStub;
+import org.tikv.kvproto.TikvGrpc.TikvFutureStub;
 
 /** Since v3.0.5 TiDB ignores the ttl on secondary lock and will use the ttl on primary key. */
 public class LockResolverClientV3 extends AbstractRegionStoreClient
@@ -79,13 +76,14 @@ public class LockResolverClientV3 extends AbstractRegionStoreClient
   public LockResolverClientV3(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
-      TikvStub asyncStub,
+      TikvFutureStub 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.RegionManager;
-import org.tikv.common.region.RegionStoreClient;
-import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.*;
 import org.tikv.common.region.TiRegion.RegionVerID;
 import org.tikv.common.util.BackOffer;
 import org.tikv.common.util.ChannelFactory;
@@ -50,7 +47,7 @@ import org.tikv.common.util.TsoUtils;
 import org.tikv.kvproto.Kvrpcpb;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
-import org.tikv.kvproto.TikvGrpc.TikvStub;
+import org.tikv.kvproto.TikvGrpc.TikvFutureStub;
 import org.tikv.txn.exception.TxnNotFoundException;
 import org.tikv.txn.exception.WriteConflictException;
 
@@ -79,13 +76,14 @@ public class LockResolverClientV4 extends AbstractRegionStoreClient
   public LockResolverClientV4(
       TiConfiguration conf,
       TiRegion region,
+      TiStore store,
       TikvBlockingStub blockingStub,
-      TikvStub asyncStub,
+      TikvFutureStub 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 =
-            this.regionManager.getRegionStorePairByKey(key, backOffer);
+        Pair<TiRegion, TiStore> pair = 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(
-      TiRegion region, Metapb.Store store, List<ByteString> keysInput, int sizeInBytes) {
+  public BatchKeys(TiRegion region, TiStore 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

@@ -1,10 +1,13 @@
 package org.tikv.common;
 
+import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
+import java.util.List;
+import java.util.stream.Collectors;
 import org.junit.Before;
-import org.tikv.common.TiConfiguration.KVMode;
 import org.tikv.common.region.TiRegion;
+import org.tikv.common.region.TiStore;
 import org.tikv.kvproto.Metapb;
 import org.tikv.kvproto.Pdpb;
 
@@ -27,14 +30,25 @@ public class MockServerTest extends PDMockServerTest {
             .addPeers(Metapb.Peer.newBuilder().setId(11).setStoreId(13))
             .build();
 
+    List<Metapb.Store> s =
+        ImmutableList.of(
+            Metapb.Store.newBuilder()
+                .setAddress("localhost:1234")
+                .setVersion("5.0.0")
+                .setId(13)
+                .build());
+
     region =
         new TiRegion(
+            session.getConf(),
             r,
             r.getPeers(0),
-            session.getConf().getIsolationLevel(),
-            session.getConf().getCommandPriority(),
-            KVMode.TXN);
+            r.getPeersList(),
+            s.stream().map(TiStore::new).collect(Collectors.toList()));
     pdServer.addGetRegionResp(Pdpb.GetRegionResponse.newBuilder().setRegion(r).build());
+    for (Metapb.Store store : s) {
+      pdServer.addGetStoreResp(Pdpb.GetStoreResponse.newBuilder().setStore(store).build());
+    }
     server = new KVMockServer();
     port = server.start(region);
   }

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

@@ -18,15 +18,14 @@ 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;
 import org.tikv.common.exception.GrpcException;
 import org.tikv.common.meta.TiTimestamp;
-import org.tikv.common.region.TiRegion;
 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.kvproto.Metapb.Store;
 import org.tikv.kvproto.Metapb.StoreState;
@@ -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(
-          ImmutableList.of("http://" + LOCAL_ADDR_IPV6 + ":" + (pdServer.port + 2)));
+      client.trySwitchLeader("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());
     }
   }
 
@@ -85,40 +85,16 @@ public class PDClientTest extends PDMockServerTest {
                 GrpcUtils.makePeer(1, 10),
                 GrpcUtils.makePeer(2, 20))));
     try (PDClient client = session.getPDClient()) {
-      TiRegion r = client.getRegionByKey(defaultBackOff(), ByteString.EMPTY);
+      Pair<Metapb.Region, Metapb.Peer> rl =
+          client.getRegionByKey(defaultBackOff(), ByteString.EMPTY);
+      Metapb.Region r = rl.first;
+      Metapb.Peer l = rl.second;
       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 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);
+      assertEquals(l.getId(), 1);
+      assertEquals(l.getStoreId(), 10);
     }
   }
 
@@ -140,40 +116,15 @@ public class PDClientTest extends PDMockServerTest {
                 GrpcUtils.makePeer(1, 10),
                 GrpcUtils.makePeer(2, 20))));
     try (PDClient client = session.getPDClient()) {
-      TiRegion r = client.getRegionByID(defaultBackOff(), 0);
+      Pair<Metapb.Region, Metapb.Peer> rl = client.getRegionByID(defaultBackOff(), 0);
+      Metapb.Region r = rl.first;
+      Metapb.Peer l = rl.second;
       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 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);
+      assertEquals(l.getId(), 1);
+      assertEquals(l.getStoreId(), 10);
     }
   }
 
@@ -208,38 +159,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 {
@@ -61,6 +61,7 @@ public class RegionManagerTest extends PDMockServerTest {
     int confVer = 1026;
     int ver = 1027;
     long regionId = 233;
+    String testAddress = "testAddress";
     pdServer.addGetRegionResp(
         GrpcUtils.makeGetRegionResponse(
             pdServer.getClusterId(),
@@ -71,6 +72,18 @@ public class RegionManagerTest extends PDMockServerTest {
                 GrpcUtils.makeRegionEpoch(confVer, ver),
                 GrpcUtils.makePeer(1, 10),
                 GrpcUtils.makePeer(2, 20))));
+    for (long id : new long[] {10, 20}) {
+      pdServer.addGetStoreResp(
+          GrpcUtils.makeGetStoreResponse(
+              pdServer.getClusterId(),
+              GrpcUtils.makeStore(
+                  id,
+                  testAddress,
+                  Metapb.StoreState.Up,
+                  GrpcUtils.makeStoreLabel("k1", "v1"),
+                  GrpcUtils.makeStoreLabel("k2", "v2"))));
+    }
+
     TiRegion region = mgr.getRegionByKey(startKey);
     assertEquals(region.getId(), regionId);
 
@@ -106,16 +119,19 @@ public class RegionManagerTest extends PDMockServerTest {
                 GrpcUtils.makeRegionEpoch(confVer, ver),
                 GrpcUtils.makePeer(storeId, 10),
                 GrpcUtils.makePeer(storeId + 1, 20))));
-    pdServer.addGetStoreResp(
-        GrpcUtils.makeGetStoreResponse(
-            pdServer.getClusterId(),
-            GrpcUtils.makeStore(
-                storeId,
-                testAddress,
-                Metapb.StoreState.Up,
-                GrpcUtils.makeStoreLabel("k1", "v1"),
-                GrpcUtils.makeStoreLabel("k2", "v2"))));
-    Pair<TiRegion, Store> pair = mgr.getRegionStorePairByKey(searchKey);
+    for (long id : new long[] {10, 20}) {
+      pdServer.addGetStoreResp(
+          GrpcUtils.makeGetStoreResponse(
+              pdServer.getClusterId(),
+              GrpcUtils.makeStore(
+                  id,
+                  testAddress,
+                  Metapb.StoreState.Up,
+                  GrpcUtils.makeStoreLabel("k1", "v1"),
+                  GrpcUtils.makeStoreLabel("k2", "v2"))));
+    }
+
+    Pair<TiRegion, TiStore> pair = mgr.getRegionStorePairByKey(searchKey);
     assertEquals(pair.first.getId(), regionId);
     assertEquals(pair.first.getId(), storeId);
   }
@@ -133,8 +149,8 @@ public class RegionManagerTest extends PDMockServerTest {
                 Metapb.StoreState.Up,
                 GrpcUtils.makeStoreLabel("k1", "v1"),
                 GrpcUtils.makeStoreLabel("k2", "v2"))));
-    Store store = mgr.getStoreById(storeId);
-    assertEquals(store.getId(), storeId);
+    TiStore store = mgr.getStoreById(storeId);
+    assertEquals(store.getStore().getId(), storeId);
 
     pdServer.addGetStoreResp(
         GrpcUtils.makeGetStoreResponse(
@@ -145,7 +161,12 @@ public class RegionManagerTest extends PDMockServerTest {
                 StoreState.Tombstone,
                 GrpcUtils.makeStoreLabel("k1", "v1"),
                 GrpcUtils.makeStoreLabel("k2", "v2"))));
-    assertNull(mgr.getStoreById(storeId + 1));
+
+    try {
+      mgr.getStoreById(storeId + 1);
+      fail();
+    } catch (Exception ignored) {
+    }
 
     mgr.invalidateStore(storeId);
     try {

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,19 +41,20 @@ 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(
             session.getConf(),
             session.getChannelFactory(),
-            new RegionManager(session.getPDClient()),
+            new RegionManager(session.getConf(), session.getPDClient()),
             session.getPDClient());
 
     return builder.build(region, store);

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

@@ -0,0 +1,40 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertEquals;
+
+import org.junit.Test;
+
+public class TiConfigurationTest {
+
+  @Test
+  public void configFileTest() {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    assertEquals("configFileTest", conf.getDBPrefix());
+  }
+
+  @Test
+  public void testGrpcIdleTimeoutValue() {
+    TiConfiguration conf = TiConfiguration.createDefault();
+    // default value
+    assertEquals(TiConfiguration.getInt(ConfigUtils.TIKV_GRPC_IDLE_TIMEOUT), conf.getIdleTimeout());
+    // new value
+    int newValue = 100000;
+    conf.setIdleTimeout(newValue);
+    assertEquals(newValue, conf.getIdleTimeout());
+  }
+}

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

@@ -0,0 +1,124 @@
+package org.tikv.common;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import com.google.protobuf.ByteString;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.After;
+import org.junit.Test;
+import org.tikv.common.region.TiRegion;
+import org.tikv.raw.RawKVClient;
+
+public class TiSessionTest {
+  private TiSession session;
+
+  @After
+  public void tearDown() throws Exception {
+    if (session != null) {
+      session.close();
+    }
+  }
+
+  @Test
+  public void closeWithRunningTaskTest() throws Exception {
+    doCloseWithRunningTaskTest(true, 0);
+  }
+
+  @Test
+  public void closeAwaitTerminationWithRunningTaskTest() throws Exception {
+    doCloseWithRunningTaskTest(false, 10000);
+  }
+
+  private void doCloseWithRunningTaskTest(boolean now, long timeoutMS) throws Exception {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    session = TiSession.create(conf);
+
+    ExecutorService executorService = session.getThreadPoolForBatchGet();
+    AtomicReference<InterruptedException> interruptedException = new AtomicReference<>();
+    executorService.submit(
+        () -> {
+          int i = 1;
+          while (true) {
+            i = i + 1;
+            try {
+              Thread.sleep(1000);
+            } catch (InterruptedException e) {
+              interruptedException.set(e);
+              break;
+            }
+          }
+        });
+
+    Thread.sleep(2000);
+
+    long startMS = System.currentTimeMillis();
+    if (now) {
+      session.close();
+      Thread.sleep(1000);
+      assertNotNull(interruptedException.get());
+      assertTrue(System.currentTimeMillis() - startMS < 2000);
+    } else {
+      session.closeAwaitTermination(timeoutMS);
+      assertNotNull(interruptedException.get());
+      assertTrue(System.currentTimeMillis() - startMS >= timeoutMS);
+    }
+  }
+
+  @Test
+  public void closeTest() throws Exception {
+    doCloseTest(true, 0);
+  }
+
+  @Test
+  public void closeAwaitTerminationTest() throws Exception {
+    doCloseTest(false, 10000);
+  }
+
+  private void doCloseTest(boolean now, long timeoutMS) throws Exception {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    session = TiSession.create(conf);
+    RawKVClient client = session.createRawClient();
+
+    // test getRegionByKey
+    ByteString key = ByteString.copyFromUtf8("key");
+    ByteString value = ByteString.copyFromUtf8("value");
+    TiRegion region = session.getRegionManager().getRegionByKey(key);
+    assertNotNull(region);
+
+    // test RawKVClient
+    client.put(key, value);
+    List<ByteString> keys = new ArrayList<>();
+    keys.add(key);
+    client.batchGet(keys);
+
+    // close TiSession
+    if (now) {
+      session.close();
+    } else {
+      session.closeAwaitTermination(timeoutMS);
+    }
+
+    // test getRegionByKey
+    try {
+      session.getRegionManager().getRegionByKey(key);
+      fail();
+    } catch (RuntimeException e) {
+      assertEquals("this TiSession is closed!", e.getMessage());
+    }
+
+    // test RawKVClient
+    try {
+      client.batchGet(keys);
+      fail();
+    } catch (RejectedExecutionException e) {
+      assertTrue(e.getMessage().contains("rejected from java.util.concurrent.ThreadPoolExecutor"));
+    }
+  }
+}

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

@@ -0,0 +1,83 @@
+package org.tikv.raw;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.protobuf.ByteString;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.After;
+import org.junit.Test;
+import org.tikv.common.TiConfiguration;
+import org.tikv.common.TiSession;
+
+public class MetricsTest {
+  private List<TiSession> sessionList = new ArrayList<>();
+
+  @After
+  public void tearDown() throws Exception {
+    for (TiSession tiSession : sessionList) {
+      if (tiSession != null) {
+        tiSession.close();
+      }
+    }
+  }
+
+  @Test
+  public void oneTiSession() throws Exception {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    conf.setMetricsEnable(true);
+    TiSession session = TiSession.create(conf);
+    sessionList.add(session);
+    RawKVClient client = session.createRawClient();
+    client.put(ByteString.copyFromUtf8("k"), ByteString.copyFromUtf8("v"));
+    ByteString result = client.get(ByteString.copyFromUtf8("k"));
+    assertEquals(result.toStringUtf8(), "v");
+    client.close();
+    session.close();
+  }
+
+  @Test
+  public void twoTiSession() throws Exception {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    conf.setMetricsEnable(true);
+
+    TiSession session1 = TiSession.create(conf);
+    sessionList.add(session1);
+    RawKVClient client1 = session1.createRawClient();
+    client1.put(ByteString.copyFromUtf8("k1"), ByteString.copyFromUtf8("v1"));
+
+    TiSession session2 = TiSession.create(conf);
+    sessionList.add(session2);
+    RawKVClient client2 = session2.createRawClient();
+    client2.put(ByteString.copyFromUtf8("k2"), ByteString.copyFromUtf8("v2"));
+
+    client1.close();
+    session1.close();
+
+    ByteString result = client2.get(ByteString.copyFromUtf8("k2"));
+    assertEquals(result.toStringUtf8(), "v2");
+
+    client2.close();
+    session2.close();
+  }
+
+  @Test
+  public void twoTiSessionWithDifferentPort() {
+    TiConfiguration conf1 = TiConfiguration.createRawDefault();
+    conf1.setMetricsEnable(true);
+    conf1.setMetricsPort(12345);
+    TiSession session1 = TiSession.create(conf1);
+    sessionList.add(session1);
+
+    TiConfiguration conf2 = TiConfiguration.createRawDefault();
+    conf2.setMetricsEnable(true);
+    conf2.setMetricsPort(54321);
+    try {
+      TiSession.create(conf2);
+      assertEquals(1, 2);
+    } catch (IllegalArgumentException e) {
+      assertEquals(
+          "Do dot support multiple tikv.metrics.port, which are 54321 and 12345", e.getMessage());
+    }
+  }
+}

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;
@@ -16,6 +19,10 @@ import org.tikv.common.TiSession;
 import org.tikv.common.codec.KeyUtils;
 import org.tikv.common.exception.TiKVException;
 import org.tikv.common.key.Key;
+import org.tikv.common.log.SlowLogEmptyImpl;
+import org.tikv.common.util.BackOffFunction;
+import org.tikv.common.util.BackOffer;
+import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.common.util.FastByteComparisons;
 import org.tikv.common.util.ScanOption;
 import org.tikv.kvproto.Kvrpcpb;
@@ -91,7 +98,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 +108,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;
@@ -155,6 +164,63 @@ public class RawKVClientTest {
         "%s%02d", RandomStringUtils.randomAlphabetic(3).toUpperCase(Locale.ROOT), r.nextInt(10000));
   }
 
+  @Test
+  public void testCustomBackOff() {
+    int timeout = 2000;
+    int sleep = 150;
+    BackOffer backOffer = ConcreteBackOffer.newCustomBackOff(timeout);
+    long s = System.currentTimeMillis();
+    try {
+      while (true) {
+        Thread.sleep(sleep);
+        backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, new Exception("t"));
+      }
+    } catch (Exception ignored) {
+    } finally {
+      long e = System.currentTimeMillis();
+      long duration = e - s;
+      assertTrue(duration >= 2900);
+    }
+  }
+
+  @Test
+  public void testDeadlineBackOff() {
+    int timeout = 2000;
+    int sleep = 150;
+    BackOffer backOffer = ConcreteBackOffer.newDeadlineBackOff(timeout, SlowLogEmptyImpl.INSTANCE);
+    long s = System.currentTimeMillis();
+    try {
+      while (true) {
+        Thread.sleep(sleep);
+        backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, new Exception("t"));
+      }
+    } catch (Exception ignored) {
+    } finally {
+      long e = System.currentTimeMillis();
+      long duration = e - s;
+      assertTrue(duration <= timeout + sleep);
+    }
+  }
+
+  @Test
+  public void testBackoffTimeout() {
+    int timeout = 500;
+    int sleep = 150;
+    BackOffer backOffer = ConcreteBackOffer.newDeadlineBackOff(timeout, SlowLogEmptyImpl.INSTANCE);
+    long s = System.currentTimeMillis();
+    try {
+      while (true) {
+        Thread.sleep(sleep);
+        backOffer.checkTimeout();
+      }
+    } catch (Exception ignored) {
+    } finally {
+      long e = System.currentTimeMillis();
+      long duration = e - s;
+      assertTrue(duration <= timeout + sleep);
+    }
+  }
+
   @Test
   public void batchPutTest() {
     if (!initialized) return;
@@ -254,30 +320,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);
-      checkEmpty(key2);
-      checkPut(key1, value1);
-      checkPut(key2, value2);
-      List<Kvrpcpb.KvPair> result = new ArrayList<>();
-      List<Kvrpcpb.KvPair> result2 = new ArrayList<>();
-      result.add(kv1);
-      result.add(kv2);
-      checkScan(key, key3, result, limit);
-      checkScan(key1, key3, result, limit);
-      checkScan(key, key1, new ArrayList<>(), limit);
-      result2.add(kv1);
-      checkScan(key, key2, result2, limit);
+      checkDeleteRange(ByteString.EMPTY, ByteString.EMPTY);
+      checkEmpty(kv);
+      checkEmpty(kv1);
+      checkEmpty(kv2);
+      checkEmpty(kv3);
+      checkPut(kv);
+      checkPut(kv1);
+      checkPut(kv2);
+      checkPut(kv3);
+      // <key, value>, <key1,value1>, <key2,value2>, <key3,value3>
+      // (-∞, +∞)
+      checkScan(ByteString.EMPTY, ByteString.EMPTY, Arrays.asList(kv, kv1, kv2, kv3), 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 +592,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 +839,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());
-      assert kvPair.getValue().equals(data.get(kvPair.getKey()));
+      assertTrue(data.containsKey(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 +899,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 +914,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;
-    assert t <= ttl && t > 0;
+    assertNotNull(t);
+    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/raw/SmartRawKVClientTest.java

@@ -0,0 +1,92 @@
+package org.tikv.raw;
+
+import static org.junit.Assert.assertTrue;
+
+import com.google.protobuf.ByteString;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.tikv.common.TiConfiguration;
+import org.tikv.common.TiSession;
+import org.tikv.common.exception.CircuitBreakerOpenException;
+
+public class SmartRawKVClientTest {
+  private boolean enable = true;
+  private int windowInSeconds = 2;
+  private int errorThresholdPercentage = 100;
+  private int requestVolumeThreshold = 10;
+  private int sleepWindowInSeconds = 1;
+  private int attemptRequestCount = 10;
+
+  private int sleepDelta = 100;
+
+  private TiSession session;
+  private SmartRawKVClient client;
+
+  @Before
+  public void setup() {
+    TiConfiguration conf = TiConfiguration.createRawDefault();
+    conf.setCircuitBreakEnable(enable);
+    conf.setCircuitBreakAvailabilityWindowInSeconds(windowInSeconds);
+    conf.setCircuitBreakAvailabilityErrorThresholdPercentage(errorThresholdPercentage);
+    conf.setCircuitBreakAvailabilityRequestVolumnThreshold(requestVolumeThreshold);
+    conf.setCircuitBreakSleepWindowInSeconds(sleepWindowInSeconds);
+    conf.setCircuitBreakAttemptRequestCount(attemptRequestCount);
+    session = TiSession.create(conf);
+    client = session.createSmartRawClient();
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    if (session != null) {
+      session.close();
+    }
+  }
+
+  @Test
+  public void testCircuitBreaker() throws InterruptedException {
+    // CLOSED => OPEN
+    {
+      for (int i = 1; i <= requestVolumeThreshold; i++) {
+        error();
+      }
+      Thread.sleep(windowInSeconds * 1000 + sleepDelta);
+
+      Exception error = null;
+      try {
+        client.get(ByteString.copyFromUtf8("key"));
+        assertTrue(false);
+      } catch (Exception e) {
+        error = e;
+      }
+      assertTrue(error instanceof CircuitBreakerOpenException);
+    }
+
+    // OPEN => CLOSED
+    {
+      Thread.sleep(sleepWindowInSeconds * 1000);
+      for (int i = 1; i <= attemptRequestCount; i++) {
+        success();
+      }
+      client.get(ByteString.copyFromUtf8("key"));
+    }
+  }
+
+  @Test
+  public void testMultiClients() {
+    for (int i = 0; i < 10240; i++) {
+      client = session.createSmartRawClient();
+    }
+  }
+
+  private void success() {
+    client.get(ByteString.copyFromUtf8("key"));
+  }
+
+  private void error() {
+    try {
+      client.callWithCircuitBreaker("error", () -> 1 / 0);
+    } catch (Exception ignored) {
+    }
+  }
+}

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 + [
-    ],
-  )