diff --git a/src/main/java/org/tikv/common/TiSession.java b/src/main/java/org/tikv/common/TiSession.java index d1171dbcf4f..6b3298ea40b 100644 --- a/src/main/java/org/tikv/common/TiSession.java +++ b/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 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 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); diff --git a/src/main/java/org/tikv/common/region/RegionCache.java b/src/main/java/org/tikv/common/region/RegionCache.java index 1b8f457da05..f9d848ac9d3 100644 --- a/src/main/java/org/tikv/common/region/RegionCache.java +++ b/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; diff --git a/src/main/java/org/tikv/common/region/RegionManager.java b/src/main/java/org/tikv/common/region/RegionManager.java index bcce5f6e94b..b9fc08d73fd 100644 --- a/src/main/java/org/tikv/common/region/RegionManager.java +++ b/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) { diff --git a/src/main/java/org/tikv/raw/SmartRawKVClient.java b/src/main/java/org/tikv/raw/SmartRawKVClient.java index a9fe1fa58a3..33a8981e81b 100644 --- a/src/main/java/org/tikv/raw/SmartRawKVClient.java +++ b/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); }