conn: fix inaccurate region count

pkg/conn/conn.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/store/tikv"
+	"github.com/pingcap/tidb/util/codec"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/backoff"
@@ -183,12 +184,17 @@ func (mgr *Mgr) GetRegionCount(ctx context.Context, startKey, endKey []byte) (in
 func (mgr *Mgr) getRegionCountWith(
 	ctx context.Context, get pdHTTPRequest, startKey, endKey []byte,
 ) (int, error) {
+	// TiKV reports region start/end keys to PD in memcomparable-format.
+	var start, end string
+	start = url.QueryEscape(string(codec.EncodeBytes(nil, startKey)))
+	if len(endKey) != 0 { // Empty end key means the max.
+		end = url.QueryEscape(string(codec.EncodeBytes(nil, endKey)))
+	}
 	var err error
 	for _, addr := range mgr.pdHTTP.addrs {
 		query := fmt.Sprintf(
 			"%s?start_key=%s&end_key=%s",
-			regionCountPrefix,
-			url.QueryEscape(string(startKey)), url.QueryEscape(string(endKey)))
+			regionCountPrefix, start, end)
 		v, e := get(ctx, addr, query, mgr.pdHTTP.cli, http.MethodGet, nil)
 		if e != nil {
 			err = e

pkg/conn/conn_test.go

@@ -2,14 +2,20 @@ package conn
 
 import (
 	"context"
+	"encoding/hex"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"testing"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/statistics"
+	"github.com/pingcap/tidb/util/codec"
 )
 
 func TestT(t *testing.T) {
@@ -22,43 +28,51 @@ type testClientSuite struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 
-	mgr *Mgr
+	mgr     *Mgr
+	regions *core.RegionsInfo
 }
 
 func (s *testClientSuite) SetUpSuite(c *C) {
 	s.ctx, s.cancel = context.WithCancel(context.Background())
 	s.mgr = &Mgr{}
+	s.regions = core.NewRegionsInfo()
 }
 
 func (s *testClientSuite) TearDownSuite(c *C) {
 	s.cancel()
 }
 
-func (s *testClientSuite) TestPDHTTP(c *C) {
-	ctx := context.Background()
+func (s *testClientSuite) TestGetClusterVersion(c *C) {
+	s.mgr.pdHTTP.addrs = []string{"", ""} // two endpoints
+	counter := 0
 	mock := func(context.Context, string, string, *http.Client, string, io.Reader) ([]byte, error) {
-		stats := statistics.RegionStats{Count: 6}
-		ret, err := json.Marshal(stats)
-		c.Assert(err, IsNil)
-		return ret, nil
-	}
-	s.mgr.pdHTTP.addrs = []string{""}
-	resp, err := s.mgr.getRegionCountWith(ctx, mock, []byte{}, []byte{})
-	c.Assert(err, IsNil)
-	c.Assert(resp, Equals, 6)
-
-	mock = func(context.Context, string, string, *http.Client, string, io.Reader) ([]byte, error) {
+		counter++
+		if counter <= 1 {
+			return nil, errors.New("mock error")
+		}
 		return []byte(`test`), nil
 	}
+
+	ctx := context.Background()
 	respString, err := s.mgr.getClusterVersionWith(ctx, mock)
 	c.Assert(err, IsNil)
 	c.Assert(respString, Equals, "test")
 
-	scheduler := "balance-leader-scheduler"
 	mock = func(context.Context, string, string, *http.Client, string, io.Reader) ([]byte, error) {
+		return nil, errors.New("mock error")
+	}
+	_, err = s.mgr.getClusterVersionWith(ctx, mock)
+	c.Assert(err, NotNil)
+}
+
+func (s *testClientSuite) TestScheduler(c *C) {
+	ctx := context.Background()
+
+	scheduler := "balance-leader-scheduler"
+	mock := func(context.Context, string, string, *http.Client, string, io.Reader) ([]byte, error) {
 		return nil, errors.New("failed")
 	}
-	err = s.mgr.removeSchedulerWith(ctx, scheduler, mock)
+	err := s.mgr.removeSchedulerWith(ctx, scheduler, mock)
 	c.Assert(err, ErrorMatches, "failed")
 
 	err = s.mgr.addSchedulerWith(ctx, scheduler, mock)
@@ -80,5 +94,56 @@ func (s *testClientSuite) TestPDHTTP(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(schedulers, HasLen, 1)
 	c.Assert(schedulers[0], Equals, scheduler)
+}
+
+func (s *testClientSuite) TestRegionCount(c *C) {
+	s.regions.SetRegion(core.NewRegionInfo(&metapb.Region{
+		Id:          1,
+		StartKey:    codec.EncodeBytes(nil, []byte{1, 1}),
+		EndKey:      codec.EncodeBytes(nil, []byte{1, 3}),
+		RegionEpoch: &metapb.RegionEpoch{},
+	}, nil))
+	s.regions.SetRegion(core.NewRegionInfo(&metapb.Region{
+		Id:          2,
+		StartKey:    codec.EncodeBytes(nil, []byte{1, 3}),
+		EndKey:      codec.EncodeBytes(nil, []byte{1, 5}),
+		RegionEpoch: &metapb.RegionEpoch{},
+	}, nil))
+	s.regions.SetRegion(core.NewRegionInfo(&metapb.Region{
+		Id:          3,
+		StartKey:    codec.EncodeBytes(nil, []byte{2, 3}),
+		EndKey:      codec.EncodeBytes(nil, []byte{3, 4}),
+		RegionEpoch: &metapb.RegionEpoch{},
+	}, nil))
+	c.Assert(s.regions.Length(), Equals, 3)
+
+	mock := func(
+		_ context.Context, addr string, prefix string, _ *http.Client, _ string, _ io.Reader,
+	) ([]byte, error) {
+		query := fmt.Sprintf("%s/%s", addr, prefix)
+		u, e := url.Parse(query)
+		c.Assert(e, IsNil, Commentf("%s", query))
+		start := u.Query().Get("start_key")
+		end := u.Query().Get("end_key")
+		c.Log(hex.EncodeToString([]byte(start)))
+		c.Log(hex.EncodeToString([]byte(end)))
+		regions := s.regions.ScanRange([]byte(start), []byte(end), 0)
+		stats := statistics.RegionStats{Count: len(regions)}
+		ret, err := json.Marshal(stats)
+		c.Assert(err, IsNil)
+		return ret, nil
+	}
+	s.mgr.pdHTTP.addrs = []string{"http://mock"}
+	ctx := context.Background()
+	resp, err := s.mgr.getRegionCountWith(ctx, mock, []byte{}, []byte{})
+	c.Assert(err, IsNil)
+	c.Assert(resp, Equals, 3)
+
+	resp, err = s.mgr.getRegionCountWith(ctx, mock, []byte{0}, []byte{0xff})
+	c.Assert(err, IsNil)
+	c.Assert(resp, Equals, 3)
 
+	resp, err = s.mgr.getRegionCountWith(ctx, mock, []byte{1, 2}, []byte{1, 4})
+	c.Assert(err, IsNil)
+	c.Assert(resp, Equals, 2)
 }