tikv: make region request can send to flash store

config/config.go

@@ -84,6 +84,7 @@ type Config struct {
 	OpenTracing         OpenTracing       `toml:"opentracing" json:"opentracing"`
 	ProxyProtocol       ProxyProtocol     `toml:"proxy-protocol" json:"proxy-protocol"`
 	TiKVClient          TiKVClient        `toml:"tikv-client" json:"tikv-client"`
+	TiFlash             TiFlash           `toml:"tiflash" json:"tiflash"`
 	Binlog              Binlog            `toml:"binlog" json:"binlog"`
 	CompatibleKillQuery bool              `toml:"compatible-kill-query" json:"compatible-kill-query"`
 	Plugin              Plugin            `toml:"plugin" json:"plugin"`
@@ -317,6 +318,13 @@ type PessimisticTxn struct {
 	TTL string `toml:"ttl" json:"ttl"`
 }
 
+// TiFlash is the config for TiFlash.
+type TiFlash struct {
+	// LabelKey and LabelValue are used to check whether a store is TiFlash.
+	LabelKey   string `toml:"label-key" json:"label-key"`
+	LabelValue string `toml:"label-value" json:"label-value"`
+}
+
 var defaultConf = Config{
 	Host:                         "0.0.0.0",
 	AdvertiseAddress:             "",
@@ -402,6 +410,10 @@ var defaultConf = Config{
 		MaxBatchWaitTime:  0,
 		BatchWaitSize:     8,
 	},
+	TiFlash: TiFlash{
+		LabelKey:   "zone",
+		LabelValue: "engine",
+	},
 	Binlog: Binlog{
 		WriteTimeout: "15s",
 		Strategy:     "range",

config/config_test.go

@@ -71,6 +71,9 @@ txn-total-size-limit=2000
 [tikv-client]
 commit-timeout="41s"
 max-batch-size=128
+[tiflash]
+label-key="zone"
+label-value="engine"
 `)
 
 	c.Assert(err, IsNil)
@@ -92,6 +95,8 @@ max-batch-size=128
 	c.Assert(conf.EnableTableLock, IsTrue)
 	c.Assert(conf.DelayCleanTableLock, Equals, uint64(5))
 	c.Assert(conf.SplitRegionMaxNum, Equals, uint64(10000))
+	c.Assert(conf.TiFlash.LabelKey, Equals, "zone")
+	c.Assert(conf.TiFlash.LabelValue, Equals, "engine")
 	c.Assert(f.Close(), IsNil)
 	c.Assert(os.Remove(configFile), IsNil)
 

store/tikv/coprocessor.go

@@ -83,9 +83,10 @@ type copTask struct {
 	region RegionVerID
 	ranges *copRanges
 
-	respChan  chan *copResponse
-	storeAddr string
-	cmdType   tikvrpc.CmdType
+	respChan    chan *copResponse
+	storeAddr   string
+	cmdType     tikvrpc.CmdType
+	isFlashTask bool
 }
 
 func (r *copTask) String() string {
@@ -624,7 +625,7 @@ func (worker *copIteratorWorker) handleTaskOnce(bo *Backoffer, task *copTask, ch
 		ScanDetail:     true,
 	})
 	startTime := time.Now()
-	resp, rpcCtx, err := sender.SendReqCtx(bo, req, task.region, ReadTimeoutMedium)
+	resp, rpcCtx, err := sender.SendReqCtx(bo, req, task.region, ReadTimeoutMedium, task.isFlashTask)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

store/tikv/region_cache.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/pd/client"
+	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/metrics"
 	"github.com/pingcap/tidb/util/logutil"
 	"go.uber.org/zap"
@@ -300,6 +301,55 @@ func (c *RegionCache) GetRPCContext(bo *Backoffer, id RegionVerID) (*RPCContext,
 	}, nil
 }
 
+// GetFlashRPCContext returns RPCContext for a region must access flash store. If it returns nil, the region
+// must be out of date and already dropped from cache or not flash store found.
+func (c *RegionCache) GetFlashRPCContext(bo *Backoffer, id RegionVerID) (*RPCContext, error) {
+	ts := time.Now().Unix()
+
+	cachedRegion := c.getCachedRegionWithRLock(id)
+	if cachedRegion == nil {
+		return nil, nil
+	}
+	if !cachedRegion.checkRegionCacheTTL(ts) {
+		return nil, nil
+	}
+
+	regionStore := cachedRegion.getStore()
+
+	for i, store := range regionStore.stores {
+		if !store.isFlash {
+			continue
+		}
+		peer, storeIdx := cachedRegion.meta.Peers[i], i
+		addr, err := c.getStoreAddr(bo, cachedRegion, store, storeIdx)
+		if err != nil {
+			return nil, err
+		}
+		if store == nil || len(addr) == 0 {
+			cachedRegion.invalidate()
+			return nil, nil
+		}
+		storeFailEpoch := atomic.LoadUint32(&store.fail)
+		if storeFailEpoch != regionStore.storeFails[regionStore.workStoreIdx] {
+			cachedRegion.invalidate()
+			logutil.BgLogger().Info("invalidate current region, because others failed on same store",
+				zap.Uint64("region", id.GetID()),
+				zap.String("store", store.addr))
+			return nil, nil
+		}
+		return &RPCContext{
+			Region:  id,
+			Meta:    cachedRegion.meta,
+			Peer:    peer,
+			PeerIdx: storeIdx,
+			Store:   store,
+			Addr:    addr,
+		}, nil
+	}
+
+	return nil, nil
+}
+
 // KeyLocation is the region and range that a key is located.
 type KeyLocation struct {
 	Region   RegionVerID
@@ -1029,6 +1079,7 @@ type Store struct {
 	state        uint64     // unsafe store storeState
 	resolveMutex sync.Mutex // protect pd from concurrent init requests
 	fail         uint32     // store fail count, see RegionStore.storeFails
+	isFlash      bool       // is the store theflash
 }
 
 type resolveState uint64
@@ -1050,6 +1101,7 @@ func (s *Store) initResolve(bo *Backoffer, c *RegionCache) (addr string, err err
 		return
 	}
 	var store *metapb.Store
+	confTiFlash := config.GetGlobalConfig().TiFlash
 	for {
 		store, err = c.pdClient.GetStore(bo.ctx, s.storeID)
 		if err != nil {
@@ -1073,6 +1125,12 @@ func (s *Store) initResolve(bo *Backoffer, c *RegionCache) (addr string, err err
 		}
 		addr = store.GetAddress()
 		s.addr = addr
+		for _, label := range store.Labels {
+			if label.Key == confTiFlash.LabelKey && label.Value == confTiFlash.LabelValue {
+				s.isFlash = true
+				break
+			}
+		}
 	retry:
 		state = s.getResolveState()
 		if state != unresolved {
@@ -1109,10 +1167,17 @@ func (s *Store) reResolve(c *RegionCache) {
 		return
 	}
 
+	confTiFlash := config.GetGlobalConfig().TiFlash
+	isFlash := false
+	for _, label := range store.Labels {
+		if label.Key == confTiFlash.LabelKey && label.Value == confTiFlash.LabelValue {
+			isFlash = true
+		}
+	}
 	addr = store.GetAddress()
 	if s.addr != addr {
 		state := resolved
-		newStore := &Store{storeID: s.storeID, addr: addr}
+		newStore := &Store{storeID: s.storeID, addr: addr, isFlash: isFlash}
 		newStore.state = *(*uint64)(unsafe.Pointer(&state))
 		c.storeMu.Lock()
 		c.storeMu.stores[newStore.storeID] = newStore

store/tikv/region_request.go

@@ -15,11 +15,12 @@ package tikv
 
 import (
 	"context"
+	"sync/atomic"
+	"time"
+
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
-	"sync/atomic"
-	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
@@ -69,12 +70,18 @@ func NewRegionRequestSender(regionCache *RegionCache, client Client) *RegionRequ
 
 // SendReq sends a request to tikv server.
 func (s *RegionRequestSender) SendReq(bo *Backoffer, req *tikvrpc.Request, regionID RegionVerID, timeout time.Duration) (*tikvrpc.Response, error) {
-	resp, _, err := s.SendReqCtx(bo, req, regionID, timeout)
+	resp, _, err := s.SendReqCtx(bo, req, regionID, timeout, false)
+	return resp, err
+}
+
+// SendReqToFlash sends a request to tiflash server.
+func (s *RegionRequestSender) SendReqToFlash(bo *Backoffer, req *tikvrpc.Request, regionID RegionVerID, timeout time.Duration) (*tikvrpc.Response, error) {
+	resp, _, err := s.SendReqCtx(bo, req, regionID, timeout, true)
 	return resp, err
 }
 
 // SendReqCtx sends a request to tikv server and return response and RPCCtx of this RPC.
-func (s *RegionRequestSender) SendReqCtx(bo *Backoffer, req *tikvrpc.Request, regionID RegionVerID, timeout time.Duration) (*tikvrpc.Response, *RPCContext, error) {
+func (s *RegionRequestSender) SendReqCtx(bo *Backoffer, req *tikvrpc.Request, regionID RegionVerID, timeout time.Duration, sendToFlash bool) (*tikvrpc.Response, *RPCContext, error) {
 	failpoint.Inject("tikvStoreSendReqResult", func(val failpoint.Value) {
 		switch val.(string) {
 		case "timeout":
@@ -95,7 +102,13 @@ func (s *RegionRequestSender) SendReqCtx(bo *Backoffer, req *tikvrpc.Request, re
 	})
 
 	for {
-		ctx, err := s.regionCache.GetRPCContext(bo, regionID)
+		var ctx *RPCContext
+		var err error
+		if sendToFlash {
+			ctx, err = s.regionCache.GetFlashRPCContext(bo, regionID)
+		} else {
+			ctx, err = s.regionCache.GetRPCContext(bo, regionID)
+		}
 		if err != nil {
 			return nil, nil, errors.Trace(err)
 		}

store/tikv/region_request_test.go

@@ -131,7 +131,7 @@ func (s *testRegionRequestSuite) TestSendReqCtx(c *C) {
 	region, err := s.cache.LocateRegionByID(s.bo, s.region)
 	c.Assert(err, IsNil)
 	c.Assert(region, NotNil)
-	resp, ctx, err := s.regionRequestSender.SendReqCtx(s.bo, req, region.Region, time.Second)
+	resp, ctx, err := s.regionRequestSender.SendReqCtx(s.bo, req, region.Region, time.Second, false)
 	c.Assert(err, IsNil)
 	c.Assert(resp.Resp, NotNil)
 	c.Assert(ctx, NotNil)