add checksum for apiv2 (#549)

* add checksum for apiv2

Signed-off-by: haojinming <jinming.hao@pingcap.com>

* fix test

Signed-off-by: haojinming <jinming.hao@pingcap.com>

* fix test

Signed-off-by: haojinming <jinming.hao@pingcap.com>

* address review comments

Signed-off-by: haojinming <jinming.hao@pingcap.com>

* address review comments

Signed-off-by: haojinming <jinming.hao@pingcap.com>

* add unit test for EncodeEncodeV2KeyRanges

Signed-off-by: haojinming <jinming.hao@pingcap.com>

Co-authored-by: iosmanthus <dengliming@pingcap.com>

integration_tests/raw/api_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"hash/crc64"
 	"io/ioutil"
 	"net/http"
 	"strings"
@@ -152,6 +153,14 @@ func (s *apiTestSuite) mustReverseScan(prefix string, start string, end string,
 	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
@@ -377,6 +386,35 @@ func (s *apiTestSuite) TestDeleteRange() {
 	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))
+	}
+	s.mustBatchPut(prefix, keys, values)
+	checksum := s.mustChecksum(prefix, "", "")
+	s.Equal(expect, checksum)
+}
+
 func (s *apiTestSuite) TearDownTest() {
 	if s.client != nil {
 		_ = s.client.Close()

internal/client/api_version.go

@@ -69,6 +69,17 @@ func EncodeV2Range(mode Mode, start, end []byte) ([]byte, []byte) {
 	return EncodeV2Key(mode, start), b
 }
 
+// EncodeV2KeyRanges encode KeyRange slice into API V2 formatted new slice.
+func EncodeV2KeyRanges(mode Mode, keyRanges []*kvrpcpb.KeyRange) []*kvrpcpb.KeyRange {
+	encodedRanges := make([]*kvrpcpb.KeyRange, 0, len(keyRanges))
+	for i := 0; i < len(keyRanges); i++ {
+		keyRange := kvrpcpb.KeyRange{}
+		keyRange.StartKey, keyRange.EndKey = EncodeV2Range(mode, keyRanges[i].StartKey, keyRanges[i].EndKey)
+		encodedRanges = append(encodedRanges, &keyRange)
+	}
+	return encodedRanges
+}
+
 // MapV2RangeToV1 maps a range in API V2 format into V1 range.
 // This function forbid the user seeing other keyspace.
 func MapV2RangeToV1(mode Mode, start []byte, end []byte) ([]byte, []byte) {
@@ -111,6 +122,7 @@ func EncodeV2Pairs(mode Mode, pairs []*kvrpcpb.KvPair) []*kvrpcpb.KvPair {
 }
 
 // EncodeRequest encodes req into specified API version format.
+// NOTE: req is reused on retry. MUST encode on cloned request, other than overwrite the original.
 func EncodeRequest(req *tikvrpc.Request) (*tikvrpc.Request, error) {
 	if req.GetApiVersion() == kvrpcpb.APIVersion_V1 {
 		return req, nil
@@ -160,6 +172,10 @@ func EncodeRequest(req *tikvrpc.Request) (*tikvrpc.Request, error) {
 		r := *req.RawCompareAndSwap()
 		r.Key = EncodeV2Key(ModeRaw, r.Key)
 		newReq.Req = &r
+	case tikvrpc.CmdRawChecksum:
+		r := *req.RawChecksum()
+		r.Ranges = EncodeV2KeyRanges(ModeRaw, r.Ranges)
+		newReq.Req = &r
 	}
 
 	return &newReq, nil

internal/client/api_version_test.go

@@ -25,3 +25,44 @@ func TestEncodeRequest(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, append(APIV2RawKeyPrefix, []byte("key")...), r.RawGet().Key)
 }
+
+func TestEncodeEncodeV2KeyRanges(t *testing.T) {
+	keyRanges := []*kvrpcpb.KeyRange{
+		{
+			StartKey: []byte{},
+			EndKey:   []byte{},
+		},
+		{
+			StartKey: []byte{},
+			EndKey:   []byte{'z'},
+		},
+		{
+			StartKey: []byte{'a'},
+			EndKey:   []byte{},
+		},
+		{
+			StartKey: []byte{'a'},
+			EndKey:   []byte{'z'},
+		},
+	}
+	expect := []*kvrpcpb.KeyRange{
+		{
+			StartKey: getV2Prefix(ModeRaw),
+			EndKey:   getV2EndKey(ModeRaw),
+		},
+		{
+			StartKey: getV2Prefix(ModeRaw),
+			EndKey:   append(getV2Prefix(ModeRaw), 'z'),
+		},
+		{
+			StartKey: append(getV2Prefix(ModeRaw), 'a'),
+			EndKey:   getV2EndKey(ModeRaw),
+		},
+		{
+			StartKey: append(getV2Prefix(ModeRaw), 'a'),
+			EndKey:   append(getV2Prefix(ModeRaw), 'z'),
+		},
+	}
+	encodedKeyRanges := EncodeV2KeyRanges(ModeRaw, keyRanges)
+	require.Equal(t, expect, encodedKeyRanges)
+}