// 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) } }