feat: support grid read write buffer

gridbuf/readbuf.go

@@ -0,0 +1,124 @@
+// Copyright 2025 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package gridbuf
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/bytedance/gopkg/lang/mcache"
+)
+
+var (
+	errReadBufferNotEnough = errors.New("error grid read buffer not enough")
+	readBufferPool         = sync.Pool{
+		New: func() interface{} {
+			return &ReadBuffer{
+				pool: make([][]byte, 0, 16),
+			}
+		},
+	}
+)
+
+type ReadBuffer struct {
+	off    int
+	chunk  []byte
+	chunks [][]byte
+	pool   [][]byte
+}
+
+func NewReadBuffer(bufs [][]byte) *ReadBuffer {
+	rb := readBufferPool.Get().(*ReadBuffer)
+	rb.chunk = bufs[0]
+	rb.chunks = bufs[1:]
+	return rb
+}
+
+// ReadN read n bytes from chunk, if chunk is not enough, it will read from next chunks.
+//
+// MAKE SURE IT CAN BE INLINE:
+// `can inline (*XReadBuffer).ReadN with cost 80`
+func (b *ReadBuffer) ReadN(n int) (buf []byte) {
+	buf = b.chunk[b.off:]
+	if len(buf) < n {
+		buf = b.readSlow(n)
+	} else {
+		b.off += n
+	}
+	return
+}
+
+func (b *ReadBuffer) readSlow(n int) (buf []byte) {
+	buf = mcache.Malloc(n)
+	b.pool = append(b.pool, buf)
+	var l, m int
+	if len(b.chunk)-b.off > 0 {
+		m = copy(buf[l:], b.chunk[b.off:])
+		l += m
+	}
+	for l < n {
+		if len(b.chunks) == 0 {
+			panic(errReadBufferNotEnough.Error())
+		}
+		b.chunk = b.chunks[0]
+		b.off = 0
+		b.chunks = b.chunks[1:]
+		m = copy(buf[l:], b.chunk)
+		l += m
+	}
+	b.off += m
+	return
+}
+
+// CopyBytes copy bytes from chunk, if chunk is not enough, it will copy from next chunks.
+//
+// MAKE SURE IT CAN BE INLINE:
+// `can inline (*XReadBuffer).CopyBytes with cost 80`
+func (b *ReadBuffer) CopyBytes(buf []byte) {
+	n := copy(buf, b.chunk[b.off:])
+	if len(buf) > n {
+		b.copySlow(buf)
+	} else {
+		b.off += n
+	}
+}
+
+func (b *ReadBuffer) copySlow(buf []byte) {
+	m := len(b.chunk) - b.off
+	l := m
+	for l < len(buf) {
+		if len(b.chunks) == 0 {
+			panic(errReadBufferNotEnough.Error())
+		}
+		b.chunk = b.chunks[0]
+		b.off = 0
+		b.chunks = b.chunks[1:]
+		m = copy(buf[l:], b.chunk)
+		l += m
+	}
+	b.off += m
+}
+
+func (b *ReadBuffer) Free() {
+	b.off = 0
+	b.chunk = nil
+	b.chunks = nil
+	for i := range b.pool {
+		mcache.Free(b.pool[i])
+		b.pool[i] = nil
+	}
+	b.pool = b.pool[:0]
+	readBufferPool.Put(b)
+}

gridbuf/readbuf_test.go

@@ -0,0 +1,171 @@
+// Copyright 2025 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package gridbuf
+
+import (
+	"runtime/debug"
+	"strings"
+	"testing"
+)
+
+func TestReadBuf_Inline(t *testing.T) {
+	var x *ReadBuffer
+
+	defer func() {
+		r := recover()
+		if r == nil {
+			t.Fatal("should panic")
+		}
+		stack := string(debug.Stack())
+		if !strings.Contains(stack, "ReadN(...)") {
+			t.Fatal("should inline ReadN")
+		}
+	}()
+	x.ReadN(10)
+
+	defer func() {
+		r := recover()
+		if r == nil {
+			t.Fatal("should panic")
+		}
+		stack := string(debug.Stack())
+		if !strings.Contains(stack, "CopyBytes(...)") {
+			t.Fatal("should inline CopyBytes")
+		}
+	}()
+
+	x.CopyBytes(make([]byte, 1))
+}
+
+func TestReadBuf_CrossPad(t *testing.T) {
+	tf := func(getBuf func(x *ReadBuffer, n int) []byte) {
+		ori1 := make([]byte, padLength)
+		for i := range ori1 {
+			ori1[i] = 'a'
+		}
+		ori2 := make([]byte, padLength)
+		for i := range ori2 {
+			ori2[i] = 'b'
+		}
+		ori := [][]byte{ori1, ori2}
+		x := NewReadBuffer(ori)
+		defer x.Free()
+		buf := getBuf(x, padLength-1)
+		if len(buf) < padLength-1 {
+			t.Fatal("buf length should be great or equal to padLength-1")
+		}
+		for _, byt := range buf[:padLength-1] {
+			if byt != 'a' {
+				t.Fatal("byt should be 'a'")
+			}
+		}
+		buf = getBuf(x, 2)
+		if len(buf) < 2 {
+			t.Fatal("buf length should be great or equal to 2")
+		}
+		if buf[0] != 'a' {
+			t.Fatal("buf[0] should be 'a'")
+		}
+		if buf[1] != 'b' {
+			t.Fatal("buf[1] should be 'b'")
+		}
+		buf = getBuf(x, padLength-1)
+		if len(buf) < padLength-1 {
+			t.Fatal("buf length should be great or equal to padLength-1")
+		}
+		for _, byt := range buf[:padLength-1] {
+			if byt != 'b' {
+				t.Fatal("byt should be 'b'")
+			}
+		}
+	}
+	tf(func(x *ReadBuffer, n int) []byte {
+		return x.ReadN(n)
+	})
+	tf(func(x *ReadBuffer, n int) []byte {
+		buf := make([]byte, n)
+		x.CopyBytes(buf)
+		return buf
+	})
+}
+
+func TestReadBuf_NoCrossPad(t *testing.T) {
+	tf := func(getBuf func(x *ReadBuffer, n int) []byte) {
+		ori1 := make([]byte, padLength/2)
+		for i := range ori1 {
+			ori1[i] = 'a'
+		}
+		ori2 := make([]byte, padLength/2)
+		for i := range ori2 {
+			ori2[i] = 'b'
+		}
+		ori := append(ori1, ori2...)
+		x := NewReadBuffer([][]byte{ori})
+		defer x.Free()
+
+		buf := getBuf(x, padLength/2)
+		if len(buf) < padLength/2 {
+			t.Fatal("buf length should be great or equal to padLength/2")
+		}
+		for _, byt := range buf[:padLength/2] {
+			if byt != 'a' {
+				t.Fatal("byt should be 'a'")
+			}
+		}
+		buf = getBuf(x, padLength/2)
+		if len(buf) < padLength/2 {
+			t.Fatal("buf length should be great or equal to padLength/2")
+		}
+		for _, byt := range buf[:padLength/2] {
+			if byt != 'b' {
+				t.Fatal("byt should be 'b'")
+			}
+		}
+	}
+	tf(func(x *ReadBuffer, n int) []byte {
+		return x.ReadN(n)
+	})
+	tf(func(x *ReadBuffer, n int) []byte {
+		buf := make([]byte, n)
+		x.CopyBytes(buf)
+		return buf
+	})
+}
+
+func BenchmarkReadBuf_ReadN(b *testing.B) {
+	bytes := make([]byte, b.N)
+	buf := NewReadBuffer([][]byte{bytes})
+	defer buf.Free()
+
+	var tmp []byte
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		tmp = buf.ReadN(1)
+	}
+	_ = tmp
+}
+
+func BenchmarkBytes_Read(b *testing.B) {
+	bytes := make([]byte, b.N)
+
+	var off int
+	var tmp []byte
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		tmp = bytes[off : off+1]
+		off++
+	}
+	_ = tmp
+}