perf: custom allocator for fast codec ReadString/ReadBinary (#1427)

pkg/protocol/bthrift/binary.go

@@ -32,19 +32,28 @@ import (
 
 var (
 	// Binary protocol for bthrift.
-	Binary          binaryProtocol
-	_               BTProtocol = binaryProtocol{}
-	spanCache                  = mem.NewSpanCache(1024 * 1024)
-	spanCacheEnable bool       = false
+	Binary binaryProtocol
+	_      BTProtocol = binaryProtocol{}
 )
 
+var allocator Allocator
+
 const binaryInplaceThreshold = 4096 // 4k
 
 type binaryProtocol struct{}
 
 // SetSpanCache enable/disable binary protocol bytes/string allocator
 func SetSpanCache(enable bool) {
-	spanCacheEnable = enable
+	if enable {
+		SetAllocator(mem.NewSpanCache(1024 * 1024))
+	} else {
+		SetAllocator(nil)
+	}
+}
+
+// SetAllocator set binary protocol bytes/string allocator.
+func SetAllocator(alloc Allocator) {
+	allocator = alloc
 }
 
 func (binaryProtocol) WriteMessageBegin(buf []byte, name string, typeID thrift.TMessageType, seqid int32) int {
@@ -473,8 +482,9 @@ func (binaryProtocol) ReadString(buf []byte) (value string, length int, err erro
 	if size < 0 || int(size) > len(buf) {
 		return value, length, perrors.NewProtocolErrorWithType(thrift.INVALID_DATA, "[ReadString] the string size greater than buf length")
 	}
-	if spanCacheEnable {
-		data := spanCache.Copy(buf[length : length+int(size)])
+	alloc := allocator
+	if alloc != nil {
+		data := alloc.Copy(buf[length : length+int(size)])
 		value = utils.SliceByteToString(data)
 	} else {
 		value = string(buf[length : length+int(size)])
@@ -494,8 +504,9 @@ func (binaryProtocol) ReadBinary(buf []byte) (value []byte, length int, err erro
 	if size < 0 || size > len(buf) {
 		return value, length, perrors.NewProtocolErrorWithType(thrift.INVALID_DATA, "[ReadBinary] the binary size greater than buf length")
 	}
-	if spanCacheEnable {
-		value = spanCache.Copy(buf[length : length+size])
+	alloc := allocator
+	if alloc != nil {
+		value = alloc.Copy(buf[length : length+size])
 	} else {
 		value = make([]byte, size)
 		copy(value, buf[length:length+size])

pkg/protocol/bthrift/binary_test.go

@@ -295,6 +295,24 @@ func TestWriteAndReadString(t *testing.T) {
 	test.Assert(t, v == "kitex")
 }
 
+// TestWriteAndReadStringWithSpanCache test binary WriteString and ReadString with spanCache allocator
+func TestWriteAndReadStringWithSpanCache(t *testing.T) {
+	buf := make([]byte, 128)
+	exceptWs := "000000056b69746578"
+	exceptSize := 9
+	wn := Binary.WriteString(buf, "kitex")
+	ws := fmt.Sprintf("%x", buf[:wn])
+	test.Assert(t, wn == exceptSize, wn, exceptSize)
+	test.Assert(t, ws == exceptWs, ws, exceptWs)
+
+	SetSpanCache(true)
+	v, length, err := Binary.ReadString(buf)
+	test.Assert(t, nil == err)
+	test.Assert(t, exceptSize == length)
+	test.Assert(t, v == "kitex")
+	SetSpanCache(false)
+}
+
 // TestWriteAndReadBinary test binary WriteBinary and ReadBinary
 func TestWriteAndReadBinary(t *testing.T) {
 	buf := make([]byte, 128)
@@ -314,6 +332,27 @@ func TestWriteAndReadBinary(t *testing.T) {
 	}
 }
 
+// TestWriteAndReadBinaryWithSpanCache test binary WriteBinary and ReadBinary with spanCache allocator
+func TestWriteAndReadBinaryWithSpanCache(t *testing.T) {
+	buf := make([]byte, 128)
+	exceptWs := "000000056b69746578"
+	exceptSize := 9
+	val := []byte("kitex")
+	wn := Binary.WriteBinary(buf, val)
+	ws := fmt.Sprintf("%x", buf[:wn])
+	test.Assert(t, wn == exceptSize, wn, exceptSize)
+	test.Assert(t, ws == exceptWs, ws, exceptWs)
+
+	SetSpanCache(true)
+	v, length, err := Binary.ReadBinary(buf)
+	test.Assert(t, nil == err)
+	test.Assert(t, exceptSize == length)
+	for i := 0; i < len(v); i++ {
+		test.Assert(t, val[i] == v[i])
+	}
+	SetSpanCache(false)
+}
+
 // TestWriteStringNocopy test binary WriteStringNocopy with small content
 func TestWriteStringNocopy(t *testing.T) {
 	buf := make([]byte, 128)

pkg/protocol/bthrift/interface.go

@@ -96,3 +96,8 @@ type BTProtocol interface {
 	ReadBinary(buf []byte) (value []byte, length int, err error)
 	Skip(buf []byte, fieldType thrift.TType) (length int, err error)
 }
+
+type Allocator interface {
+	Make(n int) []byte
+	Copy(buf []byte) (p []byte)
+}