diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index df17109c3d5..3ca4d2b5a84 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/src/main/java/org/tikv/common/TiSession.java @@ -83,9 +83,44 @@ 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 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); diff --git a/src/main/java/org/tikv/common/region/RegionCache.java b/src/main/java/org/tikv/common/region/RegionCache.java index 540e8cbad88..80869155cad 100644 --- a/src/main/java/org/tikv/common/region/RegionCache.java +++ b/src/main/java/org/tikv/common/region/RegionCache.java @@ -129,8 +129,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; diff --git a/src/main/java/org/tikv/common/region/RegionManager.java b/src/main/java/org/tikv/common/region/RegionManager.java index efc474b7db8..24b9acabe70 100644 --- a/src/main/java/org/tikv/common/region/RegionManager.java +++ b/src/main/java/org/tikv/common/region/RegionManager.java @@ -201,6 +201,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) {