// Copyright 2016 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" "fmt" . "github.com/pingcap/check" "github.com/pingcap/tidb/store/mockstore/mocktikv" ) type testRawKVSuite struct { OneByOneSuite cluster *mocktikv.Cluster client *RawKVClient bo *Backoffer } var _ = Suite(&testRawKVSuite{}) func (s *testRawKVSuite) SetUpTest(c *C) { s.cluster = mocktikv.NewCluster() mocktikv.BootstrapWithSingleStore(s.cluster) pdClient := mocktikv.NewPDClient(s.cluster) mvccStore := mocktikv.MustNewMVCCStore() s.client = &RawKVClient{ clusterID: 0, regionCache: NewRegionCache(pdClient), pdClient: pdClient, rpcClient: mocktikv.NewRPCClient(s.cluster, mvccStore), } s.bo = NewBackofferWithVars(context.Background(), 5000, nil) } func (s *testRawKVSuite) TearDownTest(c *C) { s.client.Close() } func (s *testRawKVSuite) mustNotExist(c *C, key []byte) { v, err := s.client.Get(key) c.Assert(err, IsNil) c.Assert(v, IsNil) } func (s *testRawKVSuite) mustBatchNotExist(c *C, keys [][]byte) { values, err := s.client.BatchGet(keys) c.Assert(err, IsNil) c.Assert(values, NotNil) c.Assert(len(keys), Equals, len(values)) for _, value := range values { c.Assert([]byte{}, BytesEquals, value) } } func (s *testRawKVSuite) mustGet(c *C, key, value []byte) { v, err := s.client.Get(key) c.Assert(err, IsNil) c.Assert(v, NotNil) c.Assert(v, BytesEquals, value) } func (s *testRawKVSuite) mustBatchGet(c *C, keys, values [][]byte) { checkValues, err := s.client.BatchGet(keys) c.Assert(err, IsNil) c.Assert(checkValues, NotNil) c.Assert(len(keys), Equals, len(checkValues)) for i := range keys { c.Check(values[i], BytesEquals, checkValues[i]) } } func (s *testRawKVSuite) mustPut(c *C, key, value []byte) { err := s.client.Put(key, value) c.Assert(err, IsNil) } func (s *testRawKVSuite) mustBatchPut(c *C, keys, values [][]byte) { err := s.client.BatchPut(keys, values) c.Assert(err, IsNil) } func (s *testRawKVSuite) mustDelete(c *C, key []byte) { err := s.client.Delete(key) c.Assert(err, IsNil) } func (s *testRawKVSuite) mustBatchDelete(c *C, keys [][]byte) { err := s.client.BatchDelete(keys) c.Assert(err, IsNil) } func (s *testRawKVSuite) mustScan(c *C, startKey string, limit int, expect ...string) { keys, values, err := s.client.Scan([]byte(startKey), nil, limit) c.Assert(err, IsNil) c.Assert(len(keys)*2, Equals, len(expect)) for i := range keys { c.Assert(string(keys[i]), Equals, expect[i*2]) c.Assert(string(values[i]), Equals, expect[i*2+1]) } } func (s *testRawKVSuite) mustScanRange(c *C, startKey string, endKey string, limit int, expect ...string) { keys, values, err := s.client.Scan([]byte(startKey), []byte(endKey), limit) c.Assert(err, IsNil) c.Assert(len(keys)*2, Equals, len(expect)) for i := range keys { c.Assert(string(keys[i]), Equals, expect[i*2]) c.Assert(string(values[i]), Equals, expect[i*2+1]) } } func (s *testRawKVSuite) mustReverseScan(c *C, startKey []byte, limit int, expect ...string) { keys, values, err := s.client.ReverseScan(startKey, nil, limit) c.Assert(err, IsNil) c.Assert(len(keys)*2, Equals, len(expect)) for i := range keys { c.Assert(string(keys[i]), Equals, expect[i*2]) c.Assert(string(values[i]), Equals, expect[i*2+1]) } } func (s *testRawKVSuite) mustReverseScanRange(c *C, startKey, endKey []byte, limit int, expect ...string) { keys, values, err := s.client.ReverseScan(startKey, endKey, limit) c.Assert(err, IsNil) c.Assert(len(keys)*2, Equals, len(expect)) for i := range keys { c.Assert(string(keys[i]), Equals, expect[i*2]) c.Assert(string(values[i]), Equals, expect[i*2+1]) } } func (s *testRawKVSuite) mustDeleteRange(c *C, startKey, endKey []byte, expected map[string]string) { err := s.client.DeleteRange(startKey, endKey) c.Assert(err, IsNil) for keyStr := range expected { key := []byte(keyStr) if bytes.Compare(startKey, key) <= 0 && bytes.Compare(key, endKey) < 0 { delete(expected, keyStr) } } s.checkData(c, expected) } func (s *testRawKVSuite) checkData(c *C, expected map[string]string) { keys, values, err := s.client.Scan([]byte(""), nil, len(expected)+1) c.Assert(err, IsNil) c.Assert(len(expected), Equals, len(keys)) for i, key := range keys { c.Assert(expected[string(key)], Equals, string(values[i])) } } func (s *testRawKVSuite) split(c *C, regionKey, splitKey string) error { loc, err := s.client.regionCache.LocateKey(s.bo, []byte(regionKey)) if err != nil { return err } newRegionID, peerID := s.cluster.AllocID(), s.cluster.AllocID() s.cluster.SplitRaw(loc.Region.id, newRegionID, []byte(splitKey), []uint64{peerID}, peerID) return nil } func (s *testRawKVSuite) TestSimple(c *C) { s.mustNotExist(c, []byte("key")) s.mustPut(c, []byte("key"), []byte("value")) s.mustGet(c, []byte("key"), []byte("value")) s.mustDelete(c, []byte("key")) s.mustNotExist(c, []byte("key")) err := s.client.Put([]byte("key"), []byte("")) c.Assert(err, NotNil) } func (s *testRawKVSuite) TestRawBatch(c *C) { testNum := 0 size := 0 var testKeys [][]byte var testValues [][]byte for i := 0; size/rawBatchPutSize < 4; i++ { key := fmt.Sprint("key", i) size += len(key) testKeys = append(testKeys, []byte(key)) value := fmt.Sprint("value", i) size += len(value) testValues = append(testValues, []byte(value)) s.mustNotExist(c, []byte(key)) testNum = i } err := s.split(c, "", fmt.Sprint("key", testNum/2)) c.Assert(err, IsNil) s.mustBatchPut(c, testKeys, testValues) s.mustBatchGet(c, testKeys, testValues) s.mustBatchDelete(c, testKeys) s.mustBatchNotExist(c, testKeys) } func (s *testRawKVSuite) TestSplit(c *C) { s.mustPut(c, []byte("k1"), []byte("v1")) s.mustPut(c, []byte("k3"), []byte("v3")) err := s.split(c, "k", "k2") c.Assert(err, IsNil) s.mustGet(c, []byte("k1"), []byte("v1")) s.mustGet(c, []byte("k3"), []byte("v3")) } func (s *testRawKVSuite) TestScan(c *C) { s.mustPut(c, []byte("k1"), []byte("v1")) s.mustPut(c, []byte("k3"), []byte("v3")) s.mustPut(c, []byte("k5"), []byte("v5")) s.mustPut(c, []byte("k7"), []byte("v7")) check := func() { s.mustScan(c, "", 1, "k1", "v1") s.mustScan(c, "k1", 2, "k1", "v1", "k3", "v3") s.mustScan(c, "", 10, "k1", "v1", "k3", "v3", "k5", "v5", "k7", "v7") s.mustScan(c, "k2", 2, "k3", "v3", "k5", "v5") s.mustScan(c, "k2", 3, "k3", "v3", "k5", "v5", "k7", "v7") s.mustScanRange(c, "", "k1", 1) s.mustScanRange(c, "k1", "k3", 2, "k1", "v1") s.mustScanRange(c, "k1", "k5", 10, "k1", "v1", "k3", "v3") s.mustScanRange(c, "k1", "k5\x00", 10, "k1", "v1", "k3", "v3", "k5", "v5") s.mustScanRange(c, "k5\x00", "k5\x00\x00", 10) } check() err := s.split(c, "k", "k2") c.Assert(err, IsNil) check() err = s.split(c, "k2", "k5") c.Assert(err, IsNil) check() } func (s *testRawKVSuite) TestReverseScan(c *C) { s.mustPut(c, []byte("k1"), []byte("v1")) s.mustPut(c, []byte("k3"), []byte("v3")) s.mustPut(c, []byte("k5"), []byte("v5")) s.mustPut(c, []byte("k7"), []byte("v7")) check := func() { s.mustReverseScan(c, []byte(""), 10) s.mustReverseScan(c, []byte("z"), 1, "k7", "v7") s.mustReverseScan(c, []byte("z"), 2, "k7", "v7", "k5", "v5") s.mustReverseScan(c, []byte("z"), 10, "k7", "v7", "k5", "v5", "k3", "v3", "k1", "v1") s.mustReverseScan(c, []byte("k2"), 10, "k1", "v1") s.mustReverseScan(c, []byte("k6"), 2, "k5", "v5", "k3", "v3") s.mustReverseScan(c, []byte("k5"), 1, "k3", "v3") s.mustReverseScan(c, append([]byte("k5"), 0), 1, "k5", "v5") s.mustReverseScan(c, []byte("k6"), 3, "k5", "v5", "k3", "v3", "k1", "v1") s.mustReverseScanRange(c, []byte("z"), []byte("k3"), 10, "k7", "v7", "k5", "v5", "k3", "v3") s.mustReverseScanRange(c, []byte("k7"), append([]byte("k3"), 0), 10, "k5", "v5") } check() err := s.split(c, "k", "k2") c.Assert(err, IsNil) check() err = s.split(c, "k2", "k5") c.Assert(err, IsNil) check() } func (s *testRawKVSuite) TestDeleteRange(c *C) { // Init data testData := map[string]string{} for _, i := range []byte("abcd") { for j := byte('0'); j <= byte('9'); j++ { key := []byte{i, j} value := []byte{'v', i, j} s.mustPut(c, key, value) testData[string(key)] = string(value) } } err := s.split(c, "b", "b") c.Assert(err, IsNil) err = s.split(c, "c", "c") c.Assert(err, IsNil) err = s.split(c, "d", "d") c.Assert(err, IsNil) s.checkData(c, testData) s.mustDeleteRange(c, []byte("b"), []byte("c0"), testData) s.mustDeleteRange(c, []byte("c11"), []byte("c12"), testData) s.mustDeleteRange(c, []byte("d0"), []byte("d0"), testData) s.mustDeleteRange(c, []byte("c5"), []byte("d5"), testData) s.mustDeleteRange(c, []byte("a"), []byte("z"), testData) }