github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/milevadb-server/einsteindb/delete_range_test.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package einsteindb
    15  
    16  import (
    17  	"bytes"
    18  	"context"
    19  	"math/rand"
    20  	"sort"
    21  
    22  	. "github.com/whtcorpsinc/check"
    23  	"github.com/whtcorpsinc/milevadb/causetstore/mockstore/cluster"
    24  	"github.com/whtcorpsinc/milevadb/causetstore/mockstore/mockeinsteindb"
    25  )
    26  
    27  type testDeleteRangeSuite struct {
    28  	OneByOneSuite
    29  	cluster cluster.Cluster
    30  	causetstore   *einsteindbStore
    31  }
    32  
    33  var _ = Suite(&testDeleteRangeSuite{})
    34  
    35  func (s *testDeleteRangeSuite) SetUpTest(c *C) {
    36  	client, cluster, FIDelClient, err := mockeinsteindb.NewEinsteinDBAndFIDelClient("")
    37  	c.Assert(err, IsNil)
    38  	mockeinsteindb.BootstrapWithMultiRegions(cluster, []byte("b"), []byte("c"), []byte("d"))
    39  	s.cluster = cluster
    40  	causetstore, err := NewTestEinsteinDBStore(client, FIDelClient, nil, nil, 0)
    41  	c.Check(err, IsNil)
    42  
    43  	// TODO: make this possible
    44  	// causetstore, err := mockstore.NewMockStore(
    45  	// 	mockstore.WithStoreType(mockstore.MockEinsteinDB),
    46  	// 	mockstore.WithClusterInspector(func(c cluster.Cluster) {
    47  	// 		mockstore.BootstrapWithMultiRegions(c, []byte("b"), []byte("c"), []byte("d"))
    48  	// 		s.cluster = c
    49  	// 	}),
    50  	// )
    51  	// c.Assert(err, IsNil)
    52  
    53  	s.causetstore = causetstore.(*einsteindbStore)
    54  }
    55  
    56  func (s *testDeleteRangeSuite) TearDownTest(c *C) {
    57  	err := s.causetstore.Close()
    58  	c.Assert(err, IsNil)
    59  }
    60  
    61  func (s *testDeleteRangeSuite) checkData(c *C, expectedData map[string]string) {
    62  	txn, err := s.causetstore.Begin()
    63  	c.Assert(err, IsNil)
    64  	it, err := txn.Iter([]byte("a"), nil)
    65  	c.Assert(err, IsNil)
    66  
    67  	// Scan all data and save into a map
    68  	data := map[string]string{}
    69  	for it.Valid() {
    70  		data[string(it.Key())] = string(it.Value())
    71  		err = it.Next()
    72  		c.Assert(err, IsNil)
    73  	}
    74  	err = txn.Commit(context.Background())
    75  	c.Assert(err, IsNil)
    76  
    77  	// Print log
    78  	actualKeys := make([]string, 0, len(data))
    79  	expectedKeys := make([]string, 0, len(expectedData))
    80  	for key := range data {
    81  		actualKeys = append(actualKeys, key)
    82  	}
    83  	for key := range expectedData {
    84  		expectedKeys = append(expectedKeys, key)
    85  	}
    86  	sort.Strings(actualKeys)
    87  	sort.Strings(expectedKeys)
    88  	c.Log("Actual:   ", actualKeys)
    89  	c.Log("Expected: ", expectedKeys)
    90  
    91  	// Assert data in the causetstore is the same as expected
    92  	c.Assert(data, DeepEquals, expectedData)
    93  }
    94  
    95  func (s *testDeleteRangeSuite) deleteRange(c *C, startKey []byte, endKey []byte) int {
    96  	task := NewDeleteRangeTask(s.causetstore, startKey, endKey, 1)
    97  
    98  	err := task.InterDircute(context.Background())
    99  	c.Assert(err, IsNil)
   100  
   101  	return task.CompletedRegions()
   102  }
   103  
   104  // deleteRangeFromMap deletes all keys in a given range from a map
   105  func deleteRangeFromMap(m map[string]string, startKey []byte, endKey []byte) {
   106  	for keyStr := range m {
   107  		key := []byte(keyStr)
   108  		if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 {
   109  			delete(m, keyStr)
   110  		}
   111  	}
   112  }
   113  
   114  // mustDeleteRange does delete range on both the map and the storage, and assert they are equal after deleting
   115  func (s *testDeleteRangeSuite) mustDeleteRange(c *C, startKey []byte, endKey []byte, expected map[string]string, regions int) {
   116  	completedRegions := s.deleteRange(c, startKey, endKey)
   117  	deleteRangeFromMap(expected, startKey, endKey)
   118  	s.checkData(c, expected)
   119  	c.Assert(completedRegions, Equals, regions)
   120  }
   121  
   122  func (s *testDeleteRangeSuite) TestDeleteRange(c *C) {
   123  	// Write some key-value pairs
   124  	txn, err := s.causetstore.Begin()
   125  	c.Assert(err, IsNil)
   126  
   127  	testData := map[string]string{}
   128  
   129  	// Generate a sequence of keys and random values
   130  	for _, i := range []byte("abcd") {
   131  		for j := byte('0'); j <= byte('9'); j++ {
   132  			key := []byte{i, j}
   133  			value := []byte{byte(rand.Intn(256)), byte(rand.Intn(256))}
   134  			testData[string(key)] = string(value)
   135  			err := txn.Set(key, value)
   136  			c.Assert(err, IsNil)
   137  		}
   138  	}
   139  
   140  	err = txn.Commit(context.Background())
   141  	c.Assert(err, IsNil)
   142  
   143  	s.checkData(c, testData)
   144  
   145  	s.mustDeleteRange(c, []byte("b"), []byte("c0"), testData, 2)
   146  	s.mustDeleteRange(c, []byte("c11"), []byte("c12"), testData, 1)
   147  	s.mustDeleteRange(c, []byte("d0"), []byte("d0"), testData, 0)
   148  	s.mustDeleteRange(c, []byte("d0\x00"), []byte("d1\x00"), testData, 1)
   149  	s.mustDeleteRange(c, []byte("c5"), []byte("d5"), testData, 2)
   150  	s.mustDeleteRange(c, []byte("a"), []byte("z"), testData, 4)
   151  }