From a26dff7f43e2d909bfa68be996d0ec612dfeb24b Mon Sep 17 00:00:00 2001
From: Kyle Xiao <xiaoguoqiao@bytedance.com>
Date: Thu, 6 Feb 2025 21:57:51 +0800
Subject: [PATCH] fix(thrift): reduce skip decoder mem alloc/copy

normally the ReaderSkipDecoder will be reused,
and it's hard to find out the issue with Benchmark.

In production, GC may happen all the time which make it eaiser to expose
the issue of calling `growSlow`
---
 protocol/thrift/skipdecoder.go      | 13 +++++---
 protocol/thrift/skipdecoder_test.go | 50 ++++++++++++++++++++++++++---
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/protocol/thrift/skipdecoder.go b/protocol/thrift/skipdecoder.go
index bcd7aa8..15412c1 100644
--- a/protocol/thrift/skipdecoder.go
+++ b/protocol/thrift/skipdecoder.go
@@ -182,10 +182,15 @@ func (p *ReaderSkipDecoder) Grow(n int) {
 }
 
 func (p *ReaderSkipDecoder) growSlow(n int) {
-	// mcache will take care of the size of newb
-	newb := mcache.Malloc(p.n + n)
-	copy(newb, p.b[:p.n])
-	mcache.Free(p.b)
+	// mcache will take care of the new cap of newb to be power of 2
+	newb := mcache.Malloc((len(p.b) + n) | 255) // at least 255 bytes
+	newb = newb[:cap(newb)]
+	if p.n > 0 {
+		copy(newb, p.b[:p.n])
+	}
+	if p.b != nil {
+		mcache.Free(p.b)
+	}
 	p.b = newb
 }
 
diff --git a/protocol/thrift/skipdecoder_test.go b/protocol/thrift/skipdecoder_test.go
index f8fe338..2d06380 100644
--- a/protocol/thrift/skipdecoder_test.go
+++ b/protocol/thrift/skipdecoder_test.go
@@ -263,15 +263,57 @@ func BenchmarkSkipDecoder(b *testing.B) {
 	}
 	sr.Release()
 
-	b.ResetTimer()
-	b.RunParallel(func(pb *testing.PB) {
-		sr := &BytesSkipDecoder{}
-		for pb.Next() {
+	b.Run("BytesSkipDecoder", func(b *testing.B) {
+		sr := NewBytesSkipDecoder(bs)
+		for i := 0; i < b.N; i++ {
+			_, err := sr.Next(STRUCT)
+			if err != nil {
+				b.Fatal(err)
+			}
 			sr.Reset(bs)
+		}
+		sr.Release()
+	})
+
+	b.Run("Parallel_BytesSkipDecoder", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				sr := NewBytesSkipDecoder(bs)
+				_, err := sr.Next(STRUCT)
+				if err != nil {
+					b.Fatal(err)
+				}
+				sr.Release()
+			}
+		})
+	})
+
+	b.Run("ReaderSkipDecoder", func(b *testing.B) {
+		r := bytes.NewReader(bs)
+		sr := NewReaderSkipDecoder(r)
+		for i := 0; i < b.N; i++ {
 			_, err := sr.Next(STRUCT)
 			if err != nil {
 				b.Fatal(err)
 			}
+			r.Reset(bs)
 		}
+		sr.Release()
 	})
+
+	b.Run("Parallel_ReaderSkipDecoder", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			r := bytes.NewReader(bs)
+			for pb.Next() {
+				sr := NewReaderSkipDecoder(r)
+				_, err := sr.Next(STRUCT)
+				if err != nil {
+					b.Fatal(err)
+				}
+				sr.Release()
+				r.Reset(bs)
+			}
+		})
+	})
+
 }