[close #375] warm up RawKVClient while creating it (#367) (#381)

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

@@ -83,9 +83,44 @@ public class TiSession implements AutoCloseable {
     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;
+      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);

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

@@ -129,8 +129,12 @@ public class RegionCache {
     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

@@ -201,6 +201,8 @@ public class RegionManager {
       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) {