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.
 
 

312 lines
8.9 KiB

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