From 264e943e4d2d21ee16c248d3eee51c3b38efcfc9 Mon Sep 17 00:00:00 2001
From: lei yu <leiysky@outlook.com>
Date: Wed, 22 Apr 2020 17:48:33 +0800
Subject: [PATCH] infoschema: add TiFlash to `cluster_info` table (#16684)

---
 infoschema/perfschema/tables.go |  2 +-
 infoschema/tables.go            | 26 +++++++++++++++++++++-----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/infoschema/perfschema/tables.go b/infoschema/perfschema/tables.go
index 28dbfa674e027..ff12ecb02a0df 100644
--- a/infoschema/perfschema/tables.go
+++ b/infoschema/perfschema/tables.go
@@ -265,7 +265,7 @@ func dataForRemoteProfile(ctx sessionctx.Context, nodeType, uri string, isGorout
 	)
 	switch nodeType {
 	case "tikv":
-		servers, err = infoschema.GetTiKVServerInfo(ctx)
+		servers, err = infoschema.GetStoreServerInfo(ctx)
 	case "pd":
 		servers, err = infoschema.GetPDServerInfo(ctx)
 	default:
diff --git a/infoschema/tables.go b/infoschema/tables.go
index 1adf25bdaba9a..dc0e09b4fe176 100644
--- a/infoschema/tables.go
+++ b/infoschema/tables.go
@@ -25,6 +25,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/parser/charset"
 	"github.com/pingcap/parser/model"
 	"github.com/pingcap/parser/mysql"
@@ -1166,7 +1167,7 @@ func GetClusterServerInfo(ctx sessionctx.Context) ([]ServerInfo, error) {
 
 	type retriever func(ctx sessionctx.Context) ([]ServerInfo, error)
 	var servers []ServerInfo
-	for _, r := range []retriever{GetTiDBServerInfo, GetPDServerInfo, GetTiKVServerInfo} {
+	for _, r := range []retriever{GetTiDBServerInfo, GetPDServerInfo, GetStoreServerInfo} {
 		nodes, err := r(ctx)
 		if err != nil {
 			return nil, err
@@ -1287,13 +1288,23 @@ func GetPDServerInfo(ctx sessionctx.Context) ([]ServerInfo, error) {
 	return servers, nil
 }
 
-// GetTiKVServerInfo returns all TiKV nodes information of cluster
-func GetTiKVServerInfo(ctx sessionctx.Context) ([]ServerInfo, error) {
+// GetStoreServerInfo returns all store nodes(TiKV or TiFlash) cluster information
+func GetStoreServerInfo(ctx sessionctx.Context) ([]ServerInfo, error) {
+	isTiFlashStore := func(store *metapb.Store) bool {
+		isTiFlash := false
+		for _, label := range store.Labels {
+			if label.GetKey() == "engine" && label.GetValue() == "tiflash" {
+				isTiFlash = true
+			}
+		}
+		return isTiFlash
+	}
+
 	store := ctx.GetStore()
 	// Get TiKV servers info.
 	tikvStore, ok := store.(tikv.Storage)
 	if !ok {
-		return nil, errors.Errorf("%T is not an TiKV store instance", store)
+		return nil, errors.Errorf("%T is not an TiKV or TiFlash store instance", store)
 	}
 	pdClient := tikvStore.GetRegionCache().PDClient()
 	if pdClient == nil {
@@ -1305,7 +1316,12 @@ func GetTiKVServerInfo(ctx sessionctx.Context) ([]ServerInfo, error) {
 	}
 	var servers []ServerInfo
 	for _, store := range stores {
-		tp := tikv.GetStoreTypeByMeta(store).Name()
+		var tp string
+		if isTiFlashStore(store) {
+			tp = "tiflash"
+		} else {
+			tp = tikv.GetStoreTypeByMeta(store).Name()
+		}
 		servers = append(servers, ServerInfo{
 			ServerType:     tp,
 			Address:        store.Address,