mirror of https://github.com/tikv/client-go.git
530 lines
13 KiB
Go
530 lines
13 KiB
Go
package raw_tikv_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"hash/crc64"
|
|
"io/ioutil"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/pingcap/kvproto/pkg/kvrpcpb"
|
|
"github.com/stretchr/testify/suite"
|
|
"github.com/tidwall/gjson"
|
|
"github.com/tikv/client-go/v2/rawkv"
|
|
"github.com/tikv/client-go/v2/tikv"
|
|
pd "github.com/tikv/pd/client"
|
|
"github.com/tikv/pd/client/pkg/caller"
|
|
)
|
|
|
|
func TestAPI(t *testing.T) {
|
|
if !*withTiKV {
|
|
t.Skip("skipping TestAPI because with-tikv is not enabled")
|
|
}
|
|
suite.Run(t, new(apiTestSuite))
|
|
}
|
|
|
|
type apiTestSuite struct {
|
|
suite.Suite
|
|
client *rawkv.Client
|
|
clientForCas *rawkv.Client
|
|
pdClient pd.Client
|
|
apiVersion kvrpcpb.APIVersion
|
|
}
|
|
|
|
func getConfig(url string) (string, error) {
|
|
transport := &http.Transport{}
|
|
client := http.Client{
|
|
Transport: transport,
|
|
}
|
|
defer transport.CloseIdleConnections()
|
|
resp, err := client.Get(url)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := ioutil.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return string(body), nil
|
|
}
|
|
|
|
func (s *apiTestSuite) getApiVersion(pdCli pd.Client) kvrpcpb.APIVersion {
|
|
stores, err := pdCli.GetAllStores(context.Background())
|
|
s.Nilf(err, "fail to get store, %v", err)
|
|
|
|
for _, store := range stores {
|
|
resp, err := getConfig(fmt.Sprintf("http://%s/config", store.StatusAddress))
|
|
s.Nilf(err, "fail to get config of TiKV store %s: %v", store.StatusAddress, err)
|
|
v := gjson.Get(resp, "storage.api-version")
|
|
if v.Type == gjson.Null || v.Uint() != 2 {
|
|
return kvrpcpb.APIVersion_V1
|
|
}
|
|
}
|
|
return kvrpcpb.APIVersion_V2
|
|
}
|
|
|
|
func (s *apiTestSuite) newRawKVClient(pdCli pd.Client, addrs []string) *rawkv.Client {
|
|
cli, err := rawkv.NewClientWithOpts(context.Background(), addrs, rawkv.WithAPIVersion(s.apiVersion))
|
|
s.Nil(err)
|
|
return cli
|
|
}
|
|
|
|
func (s *apiTestSuite) wrapPDClient(pdCli pd.Client) pd.Client {
|
|
var err error
|
|
if s.getApiVersion(pdCli) == kvrpcpb.APIVersion_V2 {
|
|
pdCli, err = tikv.NewCodecPDClientWithKeyspace(tikv.ModeRaw, pdCli, tikv.DefaultKeyspaceName)
|
|
}
|
|
s.Nil(err)
|
|
return pdCli
|
|
}
|
|
|
|
func (s *apiTestSuite) SetupTest() {
|
|
addrs := strings.Split(*pdAddrs, ",")
|
|
|
|
pdClient, err := pd.NewClient(caller.TestComponent, addrs, pd.SecurityOption{})
|
|
s.Nil(err)
|
|
s.apiVersion = s.getApiVersion(pdClient)
|
|
|
|
s.pdClient = s.wrapPDClient(pdClient)
|
|
|
|
client := s.newRawKVClient(pdClient, addrs)
|
|
s.client = client
|
|
|
|
clientForCas := s.newRawKVClient(pdClient, addrs)
|
|
clientForCas.SetAtomicForCAS(true)
|
|
s.clientForCas = clientForCas
|
|
}
|
|
|
|
func withPrefix(prefix, key string) string {
|
|
return prefix + key
|
|
}
|
|
|
|
func withPrefixes(prefix string, keys []string) []string {
|
|
var result []string
|
|
for i := range keys {
|
|
result = append(result, withPrefix(prefix, keys[i]))
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (s *apiTestSuite) cleanKeyPrefix(prefix string) {
|
|
end := append([]byte(prefix), 127)
|
|
err := s.client.DeleteRange(context.Background(), []byte(prefix), end)
|
|
s.Nil(err)
|
|
|
|
ks, _ := s.mustScan(prefix, "", "", 10240)
|
|
s.Empty(ks)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustPut(prefix string, key string, value string) {
|
|
err := s.client.Put(context.Background(), []byte(withPrefix(prefix, key)), []byte(value))
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustGet(prefix string, key string) string {
|
|
v, err := s.client.Get(context.Background(), []byte(withPrefix(prefix, key)))
|
|
s.Nil(err)
|
|
return string(v)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustDelete(prefix string, key string) {
|
|
err := s.client.Delete(context.Background(), []byte(withPrefix(prefix, key)))
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustScanBytes(prefix string, start string, end string, limit int) ([][]byte, [][]byte) {
|
|
end = withPrefix(prefix, end)
|
|
end += string([]byte{127})
|
|
ks, vs, err := s.client.Scan(context.Background(), []byte(withPrefix(prefix, start)), []byte(end), limit)
|
|
s.Nil(err)
|
|
|
|
for i := range ks {
|
|
ks[i] = bytes.TrimPrefix(ks[i], []byte(prefix))
|
|
}
|
|
return ks, vs
|
|
}
|
|
|
|
func (s *apiTestSuite) mustScan(prefix string, start string, end string, limit int) ([]string, []string) {
|
|
ks, vs := s.mustScanBytes(prefix, start, end, limit)
|
|
return toStrings(ks), toStrings(vs)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustReverseScanBytes(prefix string, start string, end string, limit int) ([][]byte, [][]byte) {
|
|
ks, vs, err := s.client.ReverseScan(context.Background(), []byte(withPrefix(prefix, start)), []byte(withPrefix(prefix, end)), limit)
|
|
s.Nil(err)
|
|
|
|
for i := range ks {
|
|
ks[i] = bytes.TrimPrefix(ks[i], []byte(prefix))
|
|
}
|
|
return ks, vs
|
|
}
|
|
|
|
func (s *apiTestSuite) mustReverseScan(prefix string, start string, end string, limit int) ([]string, []string) {
|
|
ks, vs := s.mustReverseScanBytes(prefix, start, end, limit)
|
|
return toStrings(ks), toStrings(vs)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustChecksum(prefix string, start string, end string) rawkv.RawChecksum {
|
|
end = withPrefix(prefix, end)
|
|
end += string([]byte{127})
|
|
checksum, err := s.client.Checksum(context.Background(), []byte(withPrefix(prefix, start)), []byte(end))
|
|
s.Nil(err)
|
|
return checksum
|
|
}
|
|
|
|
func (s *apiTestSuite) mustDeleteRange(prefix string, start, end string) {
|
|
if end == "" {
|
|
end = prefix
|
|
}
|
|
end += string([]byte{127})
|
|
err := s.client.DeleteRange(context.Background(), []byte(withPrefix(prefix, start)), []byte(withPrefix(prefix, end)))
|
|
s.Nil(err)
|
|
}
|
|
|
|
func toStrings(data [][]byte) []string {
|
|
var ss []string
|
|
for _, b := range data {
|
|
ss = append(ss, string(b))
|
|
}
|
|
return ss
|
|
}
|
|
|
|
func toBytes(data []string) [][]byte {
|
|
var bs [][]byte
|
|
for _, s := range data {
|
|
bs = append(bs, []byte(s))
|
|
}
|
|
return bs
|
|
}
|
|
|
|
func (s *apiTestSuite) mustBatchPut(prefix string, keys []string, values []string) {
|
|
s.Equal(len(keys), len(values))
|
|
keys = withPrefixes(prefix, keys)
|
|
err := s.client.BatchPut(context.Background(), toBytes(keys), toBytes(values))
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustBatchGetBytes(prefix string, keys []string) [][]byte {
|
|
keys = withPrefixes(prefix, keys)
|
|
vs, err := s.client.BatchGet(context.Background(), toBytes(keys))
|
|
s.Nil(err)
|
|
|
|
return vs
|
|
}
|
|
|
|
func (s *apiTestSuite) mustBatchGet(prefix string, keys []string) []string {
|
|
return toStrings(s.mustBatchGetBytes(prefix, keys))
|
|
}
|
|
|
|
func (s *apiTestSuite) mustBatchDelete(prefix string, keys []string) {
|
|
keys = withPrefixes(prefix, keys)
|
|
err := s.client.BatchDelete(context.Background(), toBytes(keys))
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustCASBytes(prefix, key string, old, new []byte) (bool, []byte) {
|
|
oldValue, success, err := s.clientForCas.CompareAndSwap(context.Background(), []byte(withPrefix(prefix, key)), old, new)
|
|
s.Nil(err)
|
|
return success, oldValue
|
|
}
|
|
|
|
// `old == ""` means "not exist"
|
|
func (s *apiTestSuite) mustCAS(prefix, key, old, new string) (bool, string) {
|
|
var oldValue []byte
|
|
if old != "" {
|
|
oldValue = []byte(old)
|
|
}
|
|
success, oldValue := s.mustCASBytes(prefix, key, oldValue, []byte(new))
|
|
return success, string(oldValue)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustPutWithTTL(prefix, key, value string, ttl uint64) {
|
|
err := s.client.PutWithTTL(context.Background(), []byte(withPrefix(prefix, key)), []byte(value), ttl)
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustGetKeyTTL(prefix, key string) *uint64 {
|
|
ttl, err := s.client.GetKeyTTL(context.Background(), []byte(withPrefix(prefix, key)))
|
|
s.Nil(err)
|
|
return ttl
|
|
}
|
|
|
|
func (s *apiTestSuite) mustNotExist(prefix string, key string) {
|
|
v, err := s.client.Get(context.Background(), []byte(withPrefix(prefix, key)))
|
|
s.Nil(err)
|
|
s.Nil(v)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustExist(prefix string, key string) {
|
|
v, err := s.client.Get(context.Background(), []byte(withPrefix(prefix, key)))
|
|
s.Nil(err)
|
|
s.NotNil(v)
|
|
}
|
|
|
|
func (s *apiTestSuite) mustSplitRegion(prefix string, splitKeys []string) {
|
|
var keys [][]byte
|
|
for i := range splitKeys {
|
|
keys = append(keys, []byte(withPrefix(prefix, splitKeys[i])))
|
|
}
|
|
_, err := s.pdClient.SplitRegions(context.Background(), keys)
|
|
if err != nil {
|
|
s.T().Fatalf("failed to split regions: %v", err)
|
|
}
|
|
s.Nil(err)
|
|
}
|
|
|
|
func (s *apiTestSuite) TestSimple() {
|
|
prefix := "test_simple"
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
s.mustNotExist(prefix, "key")
|
|
|
|
s.mustPut(prefix, "key", "value")
|
|
|
|
v := s.mustGet(prefix, "key")
|
|
s.Equal("value", v)
|
|
|
|
s.mustDelete(prefix, "key")
|
|
s.mustNotExist(prefix, "key")
|
|
}
|
|
|
|
func (s *apiTestSuite) TestScan() {
|
|
prefix := "test_scan"
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
var (
|
|
keys []string
|
|
values []string
|
|
)
|
|
for i := 0; i < 20480; i++ {
|
|
keys = append(keys, fmt.Sprintf("key@%v", i))
|
|
values = append(values, fmt.Sprintf("value@%v", i))
|
|
}
|
|
s.mustBatchPut(prefix, keys, values)
|
|
|
|
var splitKeys []string
|
|
for i := 0; i < 20480; i += 1024 {
|
|
splitKeys = append(splitKeys, fmt.Sprintf("key@%v", i))
|
|
}
|
|
|
|
s.mustSplitRegion(prefix, splitKeys)
|
|
|
|
keys, values = s.mustScan(prefix, keys[0], "", 10240)
|
|
s.Equal(10240, len(keys))
|
|
s.Equal(10240, len(values))
|
|
s.Equal(len(keys), len(values))
|
|
for i := range keys {
|
|
s.True(strings.HasPrefix(keys[i], "key@"))
|
|
s.True(strings.HasPrefix(values[i], "value@"))
|
|
}
|
|
}
|
|
|
|
func (s *apiTestSuite) TestReverseScan() {
|
|
prefix := "test_reverse_scan"
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
for i := 0; i < 10; i++ {
|
|
s.mustPut(prefix, fmt.Sprintf("key:%v", i), fmt.Sprintf("value:%v", i))
|
|
}
|
|
keys, values := s.mustReverseScan(prefix, "key:", "", 5)
|
|
for i := range keys {
|
|
s.Equal(fmt.Sprintf("key:%v", i), keys[len(keys)-1-i])
|
|
s.Equal(fmt.Sprintf("value:%v", i), values[len(keys)-1-i])
|
|
}
|
|
}
|
|
|
|
func (s *apiTestSuite) TestBatchOp() {
|
|
prefix := "test_batch_op"
|
|
ks := []string{"k1", "k2"}
|
|
s.cleanKeyPrefix(prefix)
|
|
s.mustBatchPut(prefix, ks, []string{"v1", "v2"})
|
|
vs := s.mustBatchGet(prefix, ks)
|
|
s.Equal("v1", vs[0])
|
|
s.Equal("v2", vs[1])
|
|
|
|
s.mustBatchDelete(prefix, ks)
|
|
s.mustNotExist(prefix, ks[0])
|
|
s.mustNotExist(prefix, ks[1])
|
|
}
|
|
|
|
func (s *apiTestSuite) TestCAS() {
|
|
prefix := "test_cas"
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
success, old := s.mustCAS(prefix, "key", "", "hello world")
|
|
s.True(success)
|
|
s.Equal("", old)
|
|
|
|
v := s.mustGet(prefix, "key")
|
|
s.Equal("hello world", v)
|
|
|
|
success, old = s.mustCAS(prefix, "key", "hello", "world")
|
|
s.False(success)
|
|
s.Equal("hello world", old)
|
|
|
|
v = s.mustGet(prefix, "key")
|
|
s.Equal("hello world", v)
|
|
|
|
success, old = s.mustCAS(prefix, "key", "hello world", "world")
|
|
s.True(success)
|
|
s.Equal("hello world", old)
|
|
|
|
v = s.mustGet(prefix, "key")
|
|
s.Equal("world", v)
|
|
}
|
|
|
|
func (s *apiTestSuite) TestTTL() {
|
|
prefix := "test_ttl"
|
|
|
|
var ttl uint64 = 2
|
|
s.mustPutWithTTL(prefix, "key", "value", ttl)
|
|
time.Sleep(time.Second * time.Duration(ttl/2))
|
|
|
|
rest := s.mustGetKeyTTL(prefix, "key")
|
|
s.NotNil(rest)
|
|
s.LessOrEqual(*rest, ttl/2)
|
|
|
|
time.Sleep(time.Second * time.Duration(ttl/2))
|
|
s.mustNotExist(prefix, "key")
|
|
|
|
rest = s.mustGetKeyTTL(prefix, "key")
|
|
s.Nil(rest)
|
|
}
|
|
|
|
func (s *apiTestSuite) TestDeleteRange() {
|
|
prefix := "test_delete_range"
|
|
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
var (
|
|
keys []string
|
|
values []string
|
|
)
|
|
for i := 0; i < 20480; i++ {
|
|
keys = append(keys, fmt.Sprintf("key@%v", i))
|
|
values = append(values, fmt.Sprintf("value@%v", i))
|
|
}
|
|
s.mustBatchPut(prefix, keys, values)
|
|
|
|
s.mustSplitRegion(prefix, []string{"key@4096"})
|
|
|
|
s.mustDeleteRange(prefix, "", "")
|
|
s.mustNotExist(prefix, "key@0")
|
|
s.mustNotExist(prefix, "key@1")
|
|
s.mustNotExist(prefix, "key@2")
|
|
}
|
|
|
|
func (s *apiTestSuite) TestRawChecksum() {
|
|
prefix := "test_checksum"
|
|
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
var (
|
|
keys []string
|
|
values []string
|
|
)
|
|
expect := rawkv.RawChecksum{}
|
|
digest := crc64.New(crc64.MakeTable(crc64.ECMA))
|
|
for i := 0; i < 20480; i++ {
|
|
key := fmt.Sprintf("key@%v", i)
|
|
value := fmt.Sprintf("value@%v", i)
|
|
keys = append(keys, key)
|
|
values = append(values, value)
|
|
|
|
digest.Reset()
|
|
digest.Write([]byte(key))
|
|
digest.Write([]byte(value))
|
|
expect.Crc64Xor ^= digest.Sum64()
|
|
expect.TotalKvs++
|
|
expect.TotalBytes += (uint64)(len(prefix) + len(key) + len(value))
|
|
if s.apiVersion == kvrpcpb.APIVersion_V2 {
|
|
expect.TotalBytes += 4 // 4 bytes key prefix of API v2
|
|
}
|
|
}
|
|
s.mustBatchPut(prefix, keys, values)
|
|
checksum := s.mustChecksum(prefix, "", "")
|
|
s.Equal(expect, checksum)
|
|
}
|
|
|
|
func (s *apiTestSuite) TestEmptyValue() {
|
|
prefix := "test_empty_value"
|
|
s.cleanKeyPrefix(prefix)
|
|
|
|
s.mustNotExist(prefix, "key")
|
|
|
|
verifyEmptyValue := func() {
|
|
s.mustExist(prefix, "key")
|
|
// get
|
|
v := s.mustGet(prefix, "key")
|
|
s.Empty(v)
|
|
// batch_get
|
|
vs := s.mustBatchGetBytes(prefix, []string{"key", "key1"})
|
|
s.Equal([][]byte{{}, nil}, vs)
|
|
// scan
|
|
keys, values := s.mustScanBytes(prefix, "key", "keyz", 10)
|
|
s.Equal([][]byte{[]byte("key")}, keys)
|
|
s.Equal([][]byte{{}}, values)
|
|
// reverse scan
|
|
keys, values = s.mustReverseScanBytes(prefix, "keyz", "key", 10)
|
|
s.Equal([][]byte{[]byte("key")}, keys)
|
|
s.Equal([][]byte{{}}, values)
|
|
}
|
|
|
|
verifyNotExist := func() {
|
|
s.mustNotExist(prefix, "key")
|
|
// batch_get
|
|
vs := s.mustBatchGetBytes(prefix, []string{"key", "key1"})
|
|
s.Equal([][]byte{nil, nil}, vs)
|
|
// scan
|
|
keys, values := s.mustScanBytes(prefix, "key", "keyz", 10)
|
|
s.Nil(keys)
|
|
s.Nil(values)
|
|
// reverse scan
|
|
keys, values = s.mustReverseScanBytes(prefix, "keyz", "key", 10)
|
|
s.Nil(keys)
|
|
s.Nil(values)
|
|
}
|
|
|
|
// put
|
|
s.mustPut(prefix, "key", "")
|
|
verifyEmptyValue()
|
|
|
|
// delete
|
|
s.mustDelete(prefix, "key")
|
|
verifyNotExist()
|
|
|
|
// batch_put
|
|
s.mustBatchPut(prefix, []string{"key"}, []string{""})
|
|
verifyEmptyValue()
|
|
|
|
// compare_and_swap, nil -> ""
|
|
s.mustDelete(prefix, "key")
|
|
ok, oldVal := s.mustCASBytes(prefix, "key", nil, []byte(""))
|
|
s.True(ok)
|
|
s.Nil(oldVal)
|
|
verifyEmptyValue()
|
|
|
|
// compare_and_swap, "" -> "val"
|
|
ok, oldVal = s.mustCASBytes(prefix, "key", []byte(""), []byte("val"))
|
|
s.True(ok)
|
|
s.Equal([]byte{}, oldVal)
|
|
}
|
|
|
|
func (s *apiTestSuite) TearDownTest() {
|
|
if s.client != nil {
|
|
s.Require().Nil(s.client.Close())
|
|
}
|
|
if s.clientForCas != nil {
|
|
s.Require().Nil(s.clientForCas.Close())
|
|
}
|
|
if s.pdClient != nil {
|
|
s.pdClient.Close()
|
|
}
|
|
}
|