server: GetAllStores should not return tombstone stores (#1440)

* server: add option to exclude tombstone stores

GetAllStores should not return tombstone stores in most cases. We can
add an option to support that if we needs it later.

Signed-off-by: Huachao Huang <huachao.huang@gmail.com>

client/client.go

@@ -58,7 +58,7 @@ type Client interface {
 	// GetAllStores gets all stores from pd.
 	// The store may expire later. Caller is responsible for caching and taking care
 	// of store change.
-	GetAllStores(ctx context.Context) ([]*metapb.Store, error)
+	GetAllStores(ctx context.Context, opts ...GetStoreOption) ([]*metapb.Store, error)
 	// Update GC safe point. TiKV will check it and do GC themselves if necessary.
 	// If the given safePoint is less than the current one, it will not be updated.
 	// Returns the new safePoint after updating.
@@ -67,6 +67,19 @@ type Client interface {
 	Close()
 }
 
+// GetStoreOp represents available options when getting stores.
+type GetStoreOp struct {
+	excludeTombstone bool
+}
+
+// GetStoreOption configures GetStoreOp.
+type GetStoreOption func(*GetStoreOp)
+
+// WithExcludeTombstone excludes tombstone stores from the result.
+func WithExcludeTombstone() GetStoreOption {
+	return func(op *GetStoreOp) { op.excludeTombstone = true }
+}
+
 type tsoRequest struct {
 	start    time.Time
 	ctx      context.Context
@@ -682,7 +695,13 @@ func (c *client) GetStore(ctx context.Context, storeID uint64) (*metapb.Store, e
 	return store, nil
 }
 
-func (c *client) GetAllStores(ctx context.Context) ([]*metapb.Store, error) {
+func (c *client) GetAllStores(ctx context.Context, opts ...GetStoreOption) ([]*metapb.Store, error) {
+	// Applies options
+	options := &GetStoreOp{}
+	for _, opt := range opts {
+		opt(options)
+	}
+
 	if span := opentracing.SpanFromContext(ctx); span != nil {
 		span = opentracing.StartSpan("pdclient.GetAllStores", opentracing.ChildOf(span.Context()))
 		defer span.Finish()
@@ -692,7 +711,8 @@ func (c *client) GetAllStores(ctx context.Context) ([]*metapb.Store, error) {
 
 	ctx, cancel := context.WithTimeout(ctx, pdTimeout)
 	resp, err := c.leaderClient().GetAllStores(ctx, &pdpb.GetAllStoresRequest{
-		Header: c.requestHeader(),
+		Header:                 c.requestHeader(),
+		ExcludeTombstoneStores: options.excludeTombstone,
 	})
 	cancel()
 

client/client_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gogo/protobuf/proto"
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
@@ -263,31 +264,46 @@ func (s *testClientSuite) TestGetStore(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(n, DeepEquals, store)
 
-	// Get a removed store should return error.
+	stores, err := s.client.GetAllStores(context.Background())
+	c.Assert(err, IsNil)
+	c.Assert(stores, DeepEquals, []*metapb.Store{store})
+
+	// Mark the store as offline.
 	err = cluster.RemoveStore(store.GetId())
 	c.Assert(err, IsNil)
+	offlineStore := proto.Clone(store).(*metapb.Store)
+	offlineStore.State = metapb.StoreState_Offline
 
 	// Get an offline store should be OK.
 	n, err = s.client.GetStore(context.Background(), store.GetId())
 	c.Assert(err, IsNil)
-	c.Assert(n.GetState(), Equals, metapb.StoreState_Offline)
+	c.Assert(n, DeepEquals, offlineStore)
+
+	// Should return offline stores.
+	stores, err = s.client.GetAllStores(context.Background())
+	c.Assert(err, IsNil)
+	c.Assert(stores, DeepEquals, []*metapb.Store{offlineStore})
 
+	// Mark the store as tombstone.
 	err = cluster.BuryStore(store.GetId(), true)
 	c.Assert(err, IsNil)
+	tombstoneStore := proto.Clone(store).(*metapb.Store)
+	tombstoneStore.State = metapb.StoreState_Tombstone
 
 	// Get a tombstone store should fail.
 	n, err = s.client.GetStore(context.Background(), store.GetId())
 	c.Assert(err, IsNil)
 	c.Assert(n, IsNil)
-}
 
-func (s *testClientSuite) TestGetAllStores(c *C) {
-	cluster := s.srv.GetRaftCluster()
-	c.Assert(cluster, NotNil)
+	// Should return tombstone stores.
+	stores, err = s.client.GetAllStores(context.Background())
+	c.Assert(err, IsNil)
+	c.Assert(stores, DeepEquals, []*metapb.Store{tombstoneStore})
 
-	stores, err := s.client.GetAllStores(context.Background())
+	// Should not return tombstone stores.
+	stores, err = s.client.GetAllStores(context.Background(), WithExcludeTombstone())
 	c.Assert(err, IsNil)
-	c.Assert(stores, DeepEquals, []*metapb.Store{store})
+	c.Assert(stores, IsNil)
 }
 
 func (s *testClientSuite) checkGCSafePoint(c *C, expectedSafePoint uint64) {

server/grpc_service.go

@@ -226,9 +226,21 @@ func (s *Server) GetAllStores(ctx context.Context, request *pdpb.GetAllStoresReq
 		return &pdpb.GetAllStoresResponse{Header: s.notBootstrappedHeader()}, nil
 	}
 
+	// Don't return tombstone stores.
+	var stores []*metapb.Store
+	if request.GetExcludeTombstoneStores() {
+		for _, store := range cluster.GetStores() {
+			if store.GetState() != metapb.StoreState_Tombstone {
+				stores = append(stores, store)
+			}
+		}
+	} else {
+		stores = cluster.GetStores()
+	}
+
 	return &pdpb.GetAllStoresResponse{
 		Header: s.header(),
-		Stores: cluster.GetStores(),
+		Stores: stores,
 	}, nil
 }