Support probe tiflash status (#2619)

Author: shiyuhang0

docs/userguide_3.0.md

@@ -18,7 +18,7 @@ CURRENT_DIR=`pwd`
 TISPARK_HOME="$(cd "`dirname "$0"`"/../..; pwd)"
 cd $TISPARK_HOME/tikv-client
 
-kvproto_hash=6ed99a08e262d8a32d6355dcba91cf99cb92074a
+kvproto_hash=4d69c6f95e683dfb5859277563bf896aca06ec34
 
 raft_rs_hash=b9891b673573fad77ebcf9bbe0969cf945841926
 

tikv-client/scripts/proto.sh

@@ -42,7 +42,7 @@
 @Accessors(chain = true)
 public class TiConfiguration implements Serializable {
   private static final DateTimeZone DEF_TIMEZONE = Converter.getLocalTimezone();
-  private static final int DEF_TIMEOUT = 10;
+  private static final int DEF_TIMEOUT = 3;
   private static final TimeUnit DEF_TIMEOUT_UNIT = TimeUnit.MINUTES;
   private static final int DEF_SCAN_BATCH_SIZE = 10480;
   private static final boolean DEF_IGNORE_TRUNCATE = true;

tikv-client/src/main/java/com/pingcap/tikv/TiConfiguration.java

@@ -39,8 +39,11 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -68,6 +71,9 @@ public class TiSession implements AutoCloseable {
   private boolean isClosed = false;
   private volatile TiTimestamp snapshotTimestamp;
   private volatile Catalog snapshotCatalog;
+  // storeStatusCache will be init at @see DAGIterator#isMppStoreAlive
+  private volatile Map<String, Boolean> storeStatusCache;
+  private ScheduledExecutorService storeStatusCacheExecutor;
 
   private TiSession(TiConfiguration conf) {
     this.conf = conf;
@@ -113,6 +119,10 @@ public static TiSession getInstance(TiConfiguration conf) {
     }
   }
 
+  public ChannelFactory getChannelFactory() {
+    return this.channelFactory;
+  }
+
   // if NewCollationEnabled is not set in configuration file,
   // we will set it to true when TiDB version is greater than or equal to v6.0.0.
   // Otherwise, we will set it to false
@@ -219,6 +229,28 @@ public synchronized RegionManager getRegionManager() {
     return res;
   }
 
+  public Map<String, Boolean> getStoreStatusCache() {
+    if (storeStatusCache == null) {
+      synchronized (this) {
+        if (storeStatusCache == null) {
+          storeStatusCache = new ConcurrentHashMap<>();
+          storeStatusCacheExecutor = Executors.newScheduledThreadPool(1);
+          storeStatusCacheExecutor.scheduleAtFixedRate(
+              () -> {
+                storeStatusCache.replaceAll(
+                    (k, v) ->
+                        RegionStoreClient.isMppAlive(
+                            channelFactory.getChannel(k, getPDClient().getHostMapping())));
+              },
+              0,
+              5,
+              TimeUnit.SECONDS);
+        }
+      }
+    }
+    return storeStatusCache;
+  }
+
   public ExecutorService getThreadPoolForIndexScan() {
     ExecutorService res = indexScanThreadPool;
     if (res == null) {
@@ -486,6 +518,9 @@ public synchronized void close() throws Exception {
     if (deleteRangeThreadPool != null) {
       deleteRangeThreadPool.shutdownNow();
     }
+    if (storeStatusCacheExecutor != null) {
+      storeStatusCacheExecutor.shutdownNow();
+    }
     if (client != null) {
       getPDClient().close();
     }

tikv-client/src/main/java/com/pingcap/tikv/TiSession.java

@@ -39,6 +39,7 @@
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.ExecutorCompletionService;
 import org.slf4j.Logger;
@@ -215,9 +216,18 @@ private SelectResponse process(RangeSplitter.RegionTask regionTask) {
       try {
         RegionStoreClient client =
             session.getRegionStoreClientBuilder().build(region, store, storeType);
+        // if mpp store is not alive, drop it and generate a new task.
+        if (storeType == TiStoreType.TiFlash && !isMppStoreAlive(store.getAddress())) {
+          logger.debug("Re-splitting region task due to TiFlash is unavailable");
+          remainTasks.addAll(
+              RangeSplitter.newSplitter(client.regionManager)
+                  .splitRangeByRegion(ranges, storeType));
+          continue;
+        }
         client.addResolvedLocks(startTs, resolvedLocks);
         Collection<RangeSplitter.RegionTask> tasks =
             client.coprocess(backOffer, dagRequest, region, ranges, responseQueue, startTs);
+
         if (tasks != null) {
           remainTasks.addAll(tasks);
         }
@@ -269,4 +279,20 @@ private Iterator<SelectResponse> processByStreaming(RangeSplitter.RegionTask reg
       throw new TiClientInternalException("Error Closing Store client.", e);
     }
   }
+
+  // See https://github.com/pingcap/tispark/pull/2619 for more details
+  public Boolean isMppStoreAlive(String address) {
+    try {
+      Map<String, Boolean> storeStatusCache = session.getStoreStatusCache();
+      return storeStatusCache.computeIfAbsent(
+          address,
+          key ->
+              RegionStoreClient.isMppAlive(
+                  session
+                      .getChannelFactory()
+                      .getChannel(address, session.getPDClient().getHostMapping())));
+    } catch (Exception e) {
+      throw new TiClientInternalException("Error get MppStore Status.", e);
+    }
+  }
 }

tikv-client/src/main/java/com/pingcap/tikv/operation/iterator/DAGIterator.java

@@ -33,7 +33,7 @@ public abstract class AbstractRegionStoreClient
     extends AbstractGRPCClient<TikvGrpc.TikvBlockingStub, TikvGrpc.TikvStub>
     implements RegionErrorReceiver {
 
-  protected final RegionManager regionManager;
+  public final RegionManager regionManager;
   protected TiRegion region;
 
   protected AbstractRegionStoreClient(

tikv-client/src/main/java/com/pingcap/tikv/region/AbstractRegionStoreClient.java

@@ -191,7 +191,7 @@ public void invalidateRegion(TiRegion region) {
   }
 
   public void invalidateRange(ByteString startKey, ByteString endKey) {
-    cache.invalidateRange(startKey,endKey);
+    cache.invalidateRange(startKey, endKey);
   }
 
   public static class RegionCache {
@@ -259,7 +259,8 @@ private synchronized TiRegion getRegionFromCache(Key key) {
     private synchronized void invalidateRange(ByteString startKey, ByteString endKey) {
       regionCache.remove(makeRange(startKey, endKey));
       if (logger.isDebugEnabled()) {
-        logger.debug(String.format("invalidateRange success, startKey[%s], endKey[%s]", startKey, endKey));
+        logger.debug(
+            String.format("invalidateRange success, startKey[%s], endKey[%s]", startKey, endKey));
       }
     }
 

tikv-client/src/main/java/com/pingcap/tikv/region/RegionManager.java

@@ -52,6 +52,7 @@
 import com.pingcap.tikv.util.Pair;
 import com.pingcap.tikv.util.RangeSplitter;
 import io.grpc.ManagedChannel;
+import io.grpc.stub.ClientCalls;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -62,6 +63,7 @@
 import java.util.Objects;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import org.slf4j.Logger;
@@ -106,6 +108,7 @@
 import org.tikv.kvproto.Kvrpcpb.TxnHeartBeatRequest;
 import org.tikv.kvproto.Kvrpcpb.TxnHeartBeatResponse;
 import org.tikv.kvproto.Metapb.Store;
+import org.tikv.kvproto.Mpp;
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
 import org.tikv.kvproto.TikvGrpc.TikvStub;
@@ -713,7 +716,7 @@ private List<RangeSplitter.RegionTask> handleCopResponse(
       // we need to invalidate cache when region not find
       if (regionError.hasRegionNotFound()) {
         logger.info("invalidateRange when Re-splitting region task because of region not find.");
-        this.regionManager.invalidateRange(region.getStartKey(),region.getEndKey());
+        this.regionManager.invalidateRange(region.getStartKey(), region.getEndKey());
       }
       // Split ranges
       return RangeSplitter.newSplitter(this.regionManager).splitRangeByRegion(ranges, storeType);
@@ -1202,6 +1205,21 @@ private void rawDeleteRangeHelper(RawDeleteRangeResponse resp) {
     }
   }
 
+  public static boolean isMppAlive(ManagedChannel channel) {
+    TikvGrpc.TikvBlockingStub stub =
+        TikvGrpc.newBlockingStub(channel).withDeadlineAfter(500, TimeUnit.MILLISECONDS);
+    Supplier<Mpp.IsAliveRequest> factory = () -> Mpp.IsAliveRequest.newBuilder().build();
+    try {
+      Mpp.IsAliveResponse resp =
+          ClientCalls.blockingUnaryCall(
+              stub.getChannel(), TikvGrpc.getIsAliveMethod(), stub.getCallOptions(), factory.get());
+      return resp != null && resp.getAvailable();
+    } catch (Exception e) {
+      logger.warn("Call mpp isAlive fail with Exception", e);
+      return false;
+    }
+  }
+
   public enum RequestTypes {
     REQ_TYPE_SELECT(101),
     REQ_TYPE_INDEX(102),