// Copyright 2018 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" "math/rand" "sort" . "github.com/pingcap/check" "github.com/pingcap/tidb/store/mockstore/mocktikv" ) type testDeleteRangeSuite struct { OneByOneSuite cluster *mocktikv.Cluster store *tikvStore } var _ = Suite(&testDeleteRangeSuite{}) func (s *testDeleteRangeSuite) SetUpTest(c *C) { s.cluster = mocktikv.NewCluster() mocktikv.BootstrapWithMultiRegions(s.cluster, []byte("b"), []byte("c"), []byte("d")) client, pdClient, err := mocktikv.NewTiKVAndPDClient(s.cluster, nil, "") c.Assert(err, IsNil) store, err := NewTestTiKVStore(client, pdClient, nil, nil, 0) c.Check(err, IsNil) s.store = store.(*tikvStore) } func (s *testDeleteRangeSuite) TearDownTest(c *C) { err := s.store.Close() c.Assert(err, IsNil) } func (s *testDeleteRangeSuite) checkData(c *C, expectedData map[string]string) { txn, err := s.store.Begin() c.Assert(err, IsNil) it, err := txn.Iter([]byte("a"), nil) c.Assert(err, IsNil) // Scan all data and save into a map data := map[string]string{} for it.Valid() { data[string(it.Key())] = string(it.Value()) err = it.Next() c.Assert(err, IsNil) } err = txn.Commit(context.Background()) c.Assert(err, IsNil) // Print log actualKeys := make([]string, 0, len(data)) expectedKeys := make([]string, 0, len(expectedData)) for key := range data { actualKeys = append(actualKeys, key) } for key := range expectedData { expectedKeys = append(expectedKeys, key) } sort.Strings(actualKeys) sort.Strings(expectedKeys) c.Log("Actual: ", actualKeys) c.Log("Expected: ", expectedKeys) // Assert data in the store is the same as expected c.Assert(data, DeepEquals, expectedData) } func (s *testDeleteRangeSuite) deleteRange(c *C, startKey []byte, endKey []byte) int { task := NewDeleteRangeTask(s.store, startKey, endKey, 1) err := task.Execute(context.Background()) c.Assert(err, IsNil) return task.CompletedRegions() } // deleteRangeFromMap deletes all keys in a given range from a map func deleteRangeFromMap(m map[string]string, startKey []byte, endKey []byte) { for keyStr := range m { key := []byte(keyStr) if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 { delete(m, keyStr) } } } // mustDeleteRange does delete range on both the map and the storage, and assert they are equal after deleting func (s *testDeleteRangeSuite) mustDeleteRange(c *C, startKey []byte, endKey []byte, expected map[string]string, regions int) { completedRegions := s.deleteRange(c, startKey, endKey) deleteRangeFromMap(expected, startKey, endKey) s.checkData(c, expected) c.Assert(completedRegions, Equals, regions) } func (s *testDeleteRangeSuite) TestDeleteRange(c *C) { // Write some key-value pairs txn, err := s.store.Begin() c.Assert(err, IsNil) testData := map[string]string{} // Generate a sequence of keys and random values for _, i := range []byte("abcd") { for j := byte('0'); j <= byte('9'); j++ { key := []byte{byte(i), byte(j)} value := []byte{byte(rand.Intn(256)), byte(rand.Intn(256))} testData[string(key)] = string(value) err := txn.Set(key, value) c.Assert(err, IsNil) } } err = txn.Commit(context.Background()) c.Assert(err, IsNil) s.checkData(c, testData) s.mustDeleteRange(c, []byte("b"), []byte("c0"), testData, 2) s.mustDeleteRange(c, []byte("c11"), []byte("c12"), testData, 1) s.mustDeleteRange(c, []byte("d0"), []byte("d0"), testData, 0) s.mustDeleteRange(c, []byte("d0\x00"), []byte("d1\x00"), testData, 1) s.mustDeleteRange(c, []byte("c5"), []byte("d5"), testData, 2) s.mustDeleteRange(c, []byte("a"), []byte("z"), testData, 4) }