From ded2fbf71038bc0f50637b8135e668ab05a17ac6 Mon Sep 17 00:00:00 2001
From: marsishandsome <marsishandsome@gmail.com>
Date: Fri, 23 Jul 2021 16:08:00 +0800
Subject: [PATCH] refactor TiSession.getSwitchTiKVModeClient

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 87410eab637..f8ad6eba08e 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 ChannelFactory getChannelFactory() {
   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 @@ private synchronized void cleanAfterTerminated() throws InterruptedException {
     if (catalog != null) {
       catalog.close();
     }
+
+    if(switchTiKVModeClient != null) {
+      switchTiKVModeClient.stopKeepTiKVToImportMode();
+    }
   }
 
   private List<ExecutorService> getExecutorServices() {