You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

229 lines
5.8 KiB

// Copyright 2019 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.
package tikv
import (
"bytes"
"context"
"errors"
"sort"
. "github.com/pingcap/check"
"github.com/pingcap/tidb/kv"
"github.com/pingcap/tidb/store/mockstore/mocktikv"
)
type testRangeTaskSuite struct {
OneByOneSuite
cluster *mocktikv.Cluster
store *tikvStore
testRanges []kv.KeyRange
expectedRanges [][]kv.KeyRange
}
var _ = Suite(&testRangeTaskSuite{})
func makeRange(startKey string, endKey string) kv.KeyRange {
return kv.KeyRange{
StartKey: []byte(startKey),
EndKey: []byte(endKey),
}
}
func (s *testRangeTaskSuite) SetUpTest(c *C) {
// Split the store at "a" to "z"
splitKeys := make([][]byte, 0)
for k := byte('a'); k <= byte('z'); k++ {
splitKeys = append(splitKeys, []byte{k})
}
// Calculate all region's ranges
allRegionRanges := []kv.KeyRange{makeRange("", "a")}
for i := 0; i < len(splitKeys)-1; i++ {
allRegionRanges = append(allRegionRanges, kv.KeyRange{
StartKey: splitKeys[i],
EndKey: splitKeys[i+1],
})
}
allRegionRanges = append(allRegionRanges, makeRange("z", ""))
s.cluster = mocktikv.NewCluster()
mocktikv.BootstrapWithMultiRegions(s.cluster, splitKeys...)
client, pdClient, err := mocktikv.NewTiKVAndPDClient(s.cluster, nil, "")
c.Assert(err, IsNil)
store, err := NewTestTiKVStore(client, pdClient, nil, nil, 0)
c.Assert(err, IsNil)
s.store = store.(*tikvStore)
s.testRanges = []kv.KeyRange{
makeRange("", ""),
makeRange("", "b"),
makeRange("b", ""),
makeRange("b", "x"),
makeRange("a", "d"),
makeRange("a\x00", "d\x00"),
makeRange("a\xff\xff\xff", "c\xff\xff\xff"),
makeRange("a1", "a2"),
makeRange("a", "a"),
makeRange("a3", "a3"),
}
s.expectedRanges = [][]kv.KeyRange{
allRegionRanges,
allRegionRanges[:2],
allRegionRanges[2:],
allRegionRanges[2:24],
{
makeRange("a", "b"),
makeRange("b", "c"),
makeRange("c", "d"),
},
{
makeRange("a\x00", "b"),
makeRange("b", "c"),
makeRange("c", "d"),
makeRange("d", "d\x00"),
},
{
makeRange("a\xff\xff\xff", "b"),
makeRange("b", "c"),
makeRange("c", "c\xff\xff\xff"),
},
{
makeRange("a1", "a2"),
},
{},
{},
}
}
func (s *testRangeTaskSuite) TearDownTest(c *C) {
err := s.store.Close()
c.Assert(err, IsNil)
}
func collect(c chan *kv.KeyRange) []kv.KeyRange {
c <- nil
ranges := make([]kv.KeyRange, 0)
for {
r := <-c
if r == nil {
break
}
ranges = append(ranges, *r)
}
return ranges
}
func (s *testRangeTaskSuite) checkRanges(c *C, obtained []kv.KeyRange, expected []kv.KeyRange) {
sort.Slice(obtained, func(i, j int) bool {
return bytes.Compare(obtained[i].StartKey, obtained[j].StartKey) < 0
})
c.Assert(obtained, DeepEquals, expected)
}
func batchRanges(ranges []kv.KeyRange, batchSize int) []kv.KeyRange {
result := make([]kv.KeyRange, 0, len(ranges))
for i := 0; i < len(ranges); i += batchSize {
lastRange := i + batchSize - 1
if lastRange >= len(ranges) {
lastRange = len(ranges) - 1
}
result = append(result, kv.KeyRange{
StartKey: ranges[i].StartKey,
EndKey: ranges[lastRange].EndKey,
})
}
return result
}
func (s *testRangeTaskSuite) testRangeTaskImpl(c *C, concurrency int) {
c.Logf("Test RangeTask, concurrency: %v", concurrency)
ranges := make(chan *kv.KeyRange, 100)
handler := func(ctx context.Context, r kv.KeyRange) (RangeTaskStat, error) {
ranges <- &r
stat := RangeTaskStat{
CompletedRegions: 1,
}
return stat, nil
}
runner := NewRangeTaskRunner("test-runner", s.store, concurrency, handler)
for regionsPerTask := 1; regionsPerTask <= 5; regionsPerTask++ {
for i, r := range s.testRanges {
runner.SetRegionsPerTask(regionsPerTask)
expectedRanges := batchRanges(s.expectedRanges[i], regionsPerTask)
err := runner.RunOnRange(context.Background(), r.StartKey, r.EndKey)
c.Assert(err, IsNil)
s.checkRanges(c, collect(ranges), expectedRanges)
c.Assert(runner.CompletedRegions(), Equals, len(expectedRanges))
c.Assert(runner.FailedRegions(), Equals, 0)
}
}
}
func (s *testRangeTaskSuite) TestRangeTask(c *C) {
for concurrency := 1; concurrency < 5; concurrency++ {
s.testRangeTaskImpl(c, concurrency)
}
}
func (s *testRangeTaskSuite) testRangeTaskErrorImpl(c *C, concurrency int) {
for i, r := range s.testRanges {
// Iterate all sub tasks and make it an error
subRanges := s.expectedRanges[i]
for _, subRange := range subRanges {
errKey := subRange.StartKey
c.Logf("Test RangeTask Error concurrency: %v, range: [%+q, %+q), errKey: %+q", concurrency, r.StartKey, r.EndKey, errKey)
handler := func(ctx context.Context, r kv.KeyRange) (RangeTaskStat, error) {
stat := RangeTaskStat{0, 0}
if bytes.Equal(r.StartKey, errKey) {
stat.FailedRegions++
return stat, errors.New("test error")
}
stat.CompletedRegions++
return stat, nil
}
runner := NewRangeTaskRunner("test-error-runner", s.store, concurrency, handler)
runner.SetRegionsPerTask(1)
err := runner.RunOnRange(context.Background(), r.StartKey, r.EndKey)
// RunOnRange returns no error only when all sub tasks are done successfully.
c.Assert(err, NotNil)
c.Assert(runner.CompletedRegions(), Less, len(subRanges))
c.Assert(runner.FailedRegions(), Equals, 1)
}
}
}
func (s *testRangeTaskSuite) TestRangeTaskError(c *C) {
for concurrency := 1; concurrency < 5; concurrency++ {
s.testRangeTaskErrorImpl(c, concurrency)
}
}