From a7605a17be5cdfaec88a5d0d5ea985d0210f95f7 Mon Sep 17 00:00:00 2001
From: shi yuhang <52435083+shiyuhang0@users.noreply.github.com>
Date: Thu, 25 May 2023 21:20:34 +0800
Subject: [PATCH] Fix health checking issue (#748)

Signed-off-by: shiyuhang <1136742008@qq.com>
---
 .../org/tikv/common/region/RegionManager.java |  7 ++---
 .../common/region/StoreHealthyChecker.java    | 30 +++++++++++++++++++
 .../java/org/tikv/common/region/TiStore.java  | 10 +++++++
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/tikv/common/region/RegionManager.java b/src/main/java/org/tikv/common/region/RegionManager.java
index 22d89821d8..31b629f1a0 100644
--- a/src/main/java/org/tikv/common/region/RegionManager.java
+++ b/src/main/java/org/tikv/common/region/RegionManager.java
@@ -256,11 +256,8 @@ public class RegionManager {
         if (!s.isReachable()) {
           continue;
         }
-        for (Metapb.StoreLabel label : s.getStore().getLabelsList()) {
-          if (label.getKey().equals(storeType.getLabelKey())
-              && label.getValue().equals(storeType.getLabelValue())) {
-            tiflashStores.add(s);
-          }
+        if (s.isTiFlash()) {
+          tiflashStores.add(s);
         }
       }
       // select a tiflash with Round-Robin strategy
diff --git a/src/main/java/org/tikv/common/region/StoreHealthyChecker.java b/src/main/java/org/tikv/common/region/StoreHealthyChecker.java
index 8d305649c4..3ae3f40d1f 100644
--- a/src/main/java/org/tikv/common/region/StoreHealthyChecker.java
+++ b/src/main/java/org/tikv/common/region/StoreHealthyChecker.java
@@ -20,17 +20,22 @@ import io.grpc.ManagedChannel;
 import io.grpc.health.v1.HealthCheckRequest;
 import io.grpc.health.v1.HealthCheckResponse;
 import io.grpc.health.v1.HealthGrpc;
+import io.grpc.stub.ClientCalls;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.tikv.common.ReadOnlyPDClient;
 import org.tikv.common.util.ChannelFactory;
 import org.tikv.common.util.ConcreteBackOffer;
 import org.tikv.kvproto.Metapb;
+import org.tikv.kvproto.Mpp;
+import org.tikv.kvproto.Mpp.IsAliveRequest;
+import org.tikv.kvproto.TikvGrpc;
 
 public class StoreHealthyChecker implements Runnable {
   private static final Logger logger = LoggerFactory.getLogger(StoreHealthyChecker.class);
@@ -75,6 +80,30 @@ public class StoreHealthyChecker implements Runnable {
 
   private boolean checkStoreHealth(TiStore store) {
     String addressStr = store.getStore().getAddress();
+    if (store.isTiFlash()) {
+      return checkTiFlashHealth(addressStr);
+    }
+    return checkTiKVHealth(addressStr);
+  }
+
+  private boolean checkTiFlashHealth(String addressStr) {
+    try {
+      ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
+      TikvGrpc.TikvBlockingStub stub =
+          TikvGrpc.newBlockingStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
+      Supplier<IsAliveRequest> factory = () -> Mpp.IsAliveRequest.newBuilder().build();
+      Mpp.IsAliveResponse resp =
+          ClientCalls.blockingUnaryCall(
+              stub.getChannel(), TikvGrpc.getIsAliveMethod(), stub.getCallOptions(), factory.get());
+      return resp != null && resp.getAvailable();
+    } catch (Exception e) {
+      logger.info(
+          "fail to check TiFlash health, regard as unhealthy. TiFlash address: " + addressStr, e);
+      return false;
+    }
+  }
+
+  private boolean checkTiKVHealth(String addressStr) {
     try {
       ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
       HealthGrpc.HealthBlockingStub stub =
@@ -83,6 +112,7 @@ public class StoreHealthyChecker implements Runnable {
       HealthCheckResponse resp = stub.check(req);
       return resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING;
     } catch (Exception e) {
+      logger.info("fail to check TiKV health, regard as unhealthy. TiKV address: " + addressStr, e);
       return false;
     }
   }
diff --git a/src/main/java/org/tikv/common/region/TiStore.java b/src/main/java/org/tikv/common/region/TiStore.java
index 8513e2b56e..5feaa246fe 100644
--- a/src/main/java/org/tikv/common/region/TiStore.java
+++ b/src/main/java/org/tikv/common/region/TiStore.java
@@ -105,4 +105,14 @@ public class TiStore implements Serializable {
   public long getId() {
     return this.store.getId();
   }
+
+  public boolean isTiFlash() {
+    for (Metapb.StoreLabel label : store.getLabelsList()) {
+      if (label.getKey().equals(TiStoreType.TiFlash.getLabelKey())
+          && label.getValue().equals(TiStoreType.TiFlash.getLabelValue())) {
+        return true;
+      }
+    }
+    return false;
+  }
 }