[close #375] warm up RawKVClient while creating it

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

@@ -20,10 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.google.protobuf.ByteString;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.stream.Collectors;
@@ -35,10 +32,7 @@
 import org.tikv.common.importer.SwitchTiKVModeClient;
 import org.tikv.common.key.Key;
 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.region.*;
 import org.tikv.common.util.*;
 import org.tikv.kvproto.ImportSstpb;
 import org.tikv.kvproto.Metapb;
@@ -99,9 +93,50 @@ public TiSession(TiConfiguration conf) {
     if (this.enableGrpcForward) {
       logger.info("enable grpc forward for high available");
     }
+    warmUp();
     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;
+      Optional<ByteString> prev = rawKVClient.get(exampleKey);
+      if (prev.isPresent()) {
+        rawKVClient.delete(exampleKey);
+        rawKVClient.putIfAbsent(exampleKey, prev.get());
+        rawKVClient.put(exampleKey, prev.get());
+      } else {
+        rawKVClient.putIfAbsent(exampleKey, ByteString.EMPTY);
+        rawKVClient.put(exampleKey, ByteString.EMPTY);
+        rawKVClient.delete(exampleKey);
+      }
+    } 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);

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

@@ -141,8 +141,12 @@ 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 == oldStore) {
+    if (originStore.equals(oldStore)) {
       storeCache.put(newStore.getId(), newStore);
       oldStore.markInvalid();
       return true;

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

@@ -206,6 +206,8 @@ private TiStore getStoreByIdWithBackOff(long id, BackOffer backOffer) {
       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) {

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

@@ -21,12 +21,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.TiConfiguration;
 import org.tikv.common.exception.CircuitBreakerOpenException;
-import org.tikv.common.exception.TiKVException;
 import org.tikv.common.util.Pair;
 import org.tikv.common.util.ScanOption;
 import org.tikv.kvproto.Kvrpcpb;
@@ -35,7 +33,6 @@
 
 public class SmartRawKVClient implements RawKVClientBase {
   private static final Logger logger = LoggerFactory.getLogger(SmartRawKVClient.class);
-  private static final AtomicBoolean warmed = new AtomicBoolean(false);
 
   private static final Histogram REQUEST_LATENCY =
       Histogram.build()
@@ -69,15 +66,6 @@ public class SmartRawKVClient implements RawKVClientBase {
   private final CircuitBreaker circuitBreaker;
 
   public SmartRawKVClient(RawKVClientBase client, TiConfiguration conf) {
-    // Warm up SmartRawKVClient to avoid the first slow call.
-    if (warmed.compareAndSet(false, true)) {
-      try {
-        logger.info("Warming up SmartRawKVClient");
-        client.get(ByteString.EMPTY);
-      } catch (final TiKVException ignored) {
-      }
-    }
-
     this.client = client;
     this.circuitBreaker = new CircuitBreakerImpl(conf);
   }