From 7ecacd7b2026eb7eeb90f7dafc3e16b1e6c1ab11 Mon Sep 17 00:00:00 2001
From: Liangliang Gu <marsishandsome@gmail.com>
Date: Thu, 29 Jul 2021 19:52:24 +0800
Subject: [PATCH] refactor TiSession.getSwitchTiKVModeClient (#240)

Signed-off-by: marsishandsome <marsishandsome@gmail.com>
---
 src/main/java/org/tikv/common/TiSession.java | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java
index 87410eab63..f8ad6eba08 100644
--- a/src/main/java/org/tikv/common/TiSession.java
+++ b/src/main/java/org/tikv/common/TiSession.java
@@ -71,6 +71,7 @@ public class TiSession implements AutoCloseable {
   private volatile RegionStoreClient.RegionStoreClientBuilder clientBuilder;
   private volatile ImporterStoreClient.ImporterStoreClientBuilder importerClientBuilder;
   private volatile boolean isClosed = false;
+  private volatile SwitchTiKVModeClient switchTiKVModeClient;
   private MetricsServer metricsServer;
   private static final int MAX_SPLIT_REGION_STACK_DEPTH = 6;
 
@@ -387,7 +388,16 @@ public class TiSession implements AutoCloseable {
   public SwitchTiKVModeClient getSwitchTiKVModeClient() {
     checkIsClosed();
 
-    return new SwitchTiKVModeClient(getPDClient(), getImporterRegionStoreClientBuilder());
+    SwitchTiKVModeClient res = switchTiKVModeClient;
+    if (res == null) {
+      synchronized (this) {
+        if (switchTiKVModeClient == null) {
+          switchTiKVModeClient = new SwitchTiKVModeClient(getPDClient(), getImporterRegionStoreClientBuilder());
+        }
+        res = switchTiKVModeClient;
+      }
+    }
+    return res;
   }
 
   /**
@@ -578,6 +588,10 @@ public class TiSession implements AutoCloseable {
     if (catalog != null) {
       catalog.close();
     }
+
+    if(switchTiKVModeClient != null) {
+      switchTiKVModeClient.stopKeepTiKVToImportMode();
+    }
   }
 
   private List<ExecutorService> getExecutorServices() {