perf: custom allocator for fast codec ReadString/ReadBinary

jizhuozhi · Jul 9, 2024 · c554c7d · c554c7d
1 parent 0807fd9
commit c554c7d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/pkg/protocol/bthrift/binary.go b/pkg/protocol/bthrift/binary.go
@@ -24,8 +24,8 @@ import (
 	"math"
 
 	"github.com/apache/thrift/lib/go/thrift"
-	"github.com/bytedance/gopkg/lang/dirtmake"
 
+	"github.com/cloudwego/kitex/pkg/mem"
 	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
 	"github.com/cloudwego/kitex/pkg/utils"
 )
@@ -42,8 +42,16 @@ const binaryInplaceThreshold = 4096 // 4k
 
 type binaryProtocol struct{}
 
+// SetSpanCache enable/disable binary protocol bytes/string allocator
+func SetSpanCache(enable bool) {
+	if enable {
+		SetAllocator(mem.NewSpanCache(1024 * 1024))
+	} else {
+		SetAllocator(nil)
+	}
+}
+
 // SetAllocator set binary protocol bytes/string allocator.
-// Note: This API is in the experiment and will be changed in the future.
 func SetAllocator(alloc Allocator) {
 	allocator = alloc
 }
@@ -500,7 +508,7 @@ func (binaryProtocol) ReadBinary(buf []byte) (value []byte, length int, err erro
 	if alloc != nil {
 		value = alloc.Copy(buf[length : length+size])
 	} else {
-		value = dirtmake.Bytes(size, size)
+		value = make([]byte, size)
 		copy(value, buf[length:length+size])
 	}
 	length += size

diff --git a/pkg/protocol/bthrift/binary_test.go b/pkg/protocol/bthrift/binary_test.go
@@ -22,7 +22,6 @@ import (
 	"testing"
 
 	"github.com/apache/thrift/lib/go/thrift"
-	"github.com/cloudwego/kitex/pkg/mem"
 	"github.com/cloudwego/netpoll"
 
 	"github.com/cloudwego/kitex/internal/test"
@@ -296,8 +295,8 @@ func TestWriteAndReadString(t *testing.T) {
 	test.Assert(t, v == "kitex")
 }
 
-// TestWriteAndReadStringWithAllocator test binary WriteString and ReadString with spanCache allocator
-func TestWriteAndReadStringWithAllocator(t *testing.T) {
+// TestWriteAndReadStringWithSpanCache test binary WriteString and ReadString with spanCache allocator
+func TestWriteAndReadStringWithSpanCache(t *testing.T) {
 	buf := make([]byte, 128)
 	exceptWs := "000000056b69746578"
 	exceptSize := 9
@@ -306,12 +305,12 @@ func TestWriteAndReadStringWithAllocator(t *testing.T) {
 	test.Assert(t, wn == exceptSize, wn, exceptSize)
 	test.Assert(t, ws == exceptWs, ws, exceptWs)
 
-	SetAllocator(mem.NewSpanCache(1024 * 1024))
+	SetSpanCache(true)
 	v, length, err := Binary.ReadString(buf)
 	test.Assert(t, nil == err)
 	test.Assert(t, exceptSize == length)
 	test.Assert(t, v == "kitex")
-	SetAllocator(nil)
+	SetSpanCache(false)
 }
 
 // TestWriteAndReadBinary test binary WriteBinary and ReadBinary
@@ -333,8 +332,8 @@ func TestWriteAndReadBinary(t *testing.T) {
 	}
 }
 
-// TestWriteAndReadBinaryWithAllocator test binary WriteBinary and ReadBinary with spanCache allocator
-func TestWriteAndReadBinaryWithAllocator(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
@@ -344,14 +343,14 @@ func TestWriteAndReadBinaryWithAllocator(t *testing.T) {
 	test.Assert(t, wn == exceptSize, wn, exceptSize)
 	test.Assert(t, ws == exceptWs, ws, exceptWs)
 
-	SetAllocator(mem.NewSpanCache(1024 * 1024))
+	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])
 	}
-	SetAllocator(nil)
+	SetSpanCache(false)
 }
 
 // TestWriteStringNocopy test binary WriteStringNocopy with small content