kv,store/copr: make kv.KeyRange the same with coprocessor.KeyRange, s…

…o we can use unsafe optimization (#37288)

close #37286

kv/key.go

@@ -89,9 +89,17 @@ func (k Key) String() string {
 }
 
 // KeyRange represents a range where StartKey <= key < EndKey.
+// Hack: make the layout exactly the same with github.com/pingcap/kvproto/pkg/coprocessor.KeyRange
+// So we can avoid allocation of converting kv.KeyRange to coprocessor.KeyRange
+// Not defined as "type KeyRange = coprocessor.KeyRange" because their field name are different.
+// kv.KeyRange use StartKey,EndKey while coprocessor.KeyRange use Start,End
 type KeyRange struct {
 	StartKey Key
 	EndKey   Key
+
+	XXXNoUnkeyedLiteral struct{}
+	XXXunrecognized     []byte
+	XXXsizecache        int32
 }
 
 // IsPoint checks if the key range represents a point.

kv/key_test.go

@@ -20,7 +20,9 @@ import (
 	"strconv"
 	"testing"
 	"time"
+	"unsafe"
 
+	"github.com/pingcap/kvproto/pkg/coprocessor"
 	. "github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/testkit/testutil"
@@ -212,6 +214,18 @@ func TestHandleMap(t *testing.T) {
 	assert.Equal(t, 2, cnt)
 }
 
+func TestKeyRangeDefinition(t *testing.T) {
+	// The struct layout for kv.KeyRange and coprocessor.KeyRange should be exactly the same.
+	// This allow us to use unsafe pointer to convert them and reduce allocation.
+	var r1 KeyRange
+	var r2 coprocessor.KeyRange
+	// Same size.
+	require.Equal(t, unsafe.Sizeof(r1), unsafe.Sizeof(r2))
+	// And same default value.
+	require.Equal(t, (*coprocessor.KeyRange)(unsafe.Pointer(&r1)), &r2)
+	require.Equal(t, &r1, (*KeyRange)(unsafe.Pointer(&r2)))
+}
+
 func BenchmarkIsPoint(b *testing.B) {
 	b.ReportAllocs()
 	kr := KeyRange{

store/copr/key_ranges.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"fmt"
 	"sort"
+	"unsafe"
 
 	"github.com/pingcap/kvproto/pkg/coprocessor"
 	"github.com/pingcap/tidb/kv"
@@ -104,8 +105,7 @@ func (r *KeyRanges) Do(f func(ran *kv.KeyRange)) {
 		f(r.first)
 	}
 	for i := range r.mid {
-		ran := r.mid[i]
-		f(&ran)
+		f(&r.mid[i])
 	}
 	if r.last != nil {
 		f(r.last)
@@ -136,10 +136,9 @@ func (r *KeyRanges) Split(key []byte) (*KeyRanges, *KeyRanges) {
 func (r *KeyRanges) ToPBRanges() []*coprocessor.KeyRange {
 	ranges := make([]*coprocessor.KeyRange, 0, r.Len())
 	r.Do(func(ran *kv.KeyRange) {
-		ranges = append(ranges, &coprocessor.KeyRange{
-			Start: ran.StartKey,
-			End:   ran.EndKey,
-		})
+		// kv.KeyRange and coprocessor.KeyRange are the same,
+		// so use unsafe.Pointer to avoid allocation here.
+		ranges = append(ranges, (*coprocessor.KeyRange)(unsafe.Pointer(ran)))
 	})
 	return ranges
 }

util/logutil/hex.go

@@ -57,7 +57,7 @@ func prettyPrint(w io.Writer, val reflect.Value) {
 		for i := 0; i < val.NumField(); i++ {
 			fv := val.Field(i)
 			ft := tp.Field(i)
-			if strings.HasPrefix(ft.Name, "XXX_") {
+			if strings.HasPrefix(ft.Name, "XXX") {
 				continue
 			}
 			if i != 0 {