feat: add span cache (#213)

lang/span/span.go

@@ -0,0 +1,120 @@
+// Copyright 2024 ByteDance Inc.
+//
+// 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 span
+
+import (
+	"math/bits"
+	"sync/atomic"
+
+	"github.com/bytedance/gopkg/lang/dirtmake"
+)
+
+/* Span Cache: A thread-safe linear allocator
+
+Design:
+1. [GC Friendly]: Centralize a batch of small bytes slice into a big size bytes slice to avoid malloc too many objects
+2. [Thread Safe]: Multi thread may share a same span, but it should fall back to the native allocator if lock conflict
+*/
+
+const (
+	spanCacheSize = 10
+	minSpanObject = 128                                  // 128 B
+	maxSpanObject = (minSpanObject << spanCacheSize) - 1 // 128 KB
+	minSpanClass  = 8                                    // = spanClass(minSpanObject)
+)
+
+type spanCache struct {
+	spans [spanCacheSize]*span
+}
+
+// NewSpanCache returns a spanCache with the given spanSize,
+// each span is used to allocate a binary of a specific size level.
+func NewSpanCache(spanSize int) *spanCache {
+	c := new(spanCache)
+	for i := 0; i < len(c.spans); i++ {
+		c.spans[i] = NewSpan(spanSize)
+	}
+	return c
+}
+
+// Make returns a [:n:n] bytes slice from a cached buffer
+// NOTE: Make will not clear the underlay bytes for performance concern. So caller MUST set every byte before read.
+func (c *spanCache) Make(n int) []byte {
+	sclass := spanClass(n) - minSpanClass
+	if sclass < 0 || sclass >= len(c.spans) {
+		return dirtmake.Bytes(n, n)
+	}
+	return c.spans[sclass].Make(n)
+}
+
+// Copy is an alias function for make-and-copy pattern
+func (c *spanCache) Copy(buf []byte) (p []byte) {
+	p = c.Make(len(buf))
+	copy(p, buf)
+	return p
+}
+
+// NewSpan returns a span with given size
+func NewSpan(size int) *span {
+	sp := new(span)
+	sp.size = uint32(size)
+	sp.buffer = dirtmake.Bytes(0, size)
+	return sp
+}
+
+type span struct {
+	lock   uint32
+	read   uint32 // read index of buffer
+	size   uint32 // size of buffer
+	buffer []byte
+}
+
+// Make returns a [:n:n] bytes slice from a cached buffer
+// NOTE: Make will not clear the underlay bytes for performance concern. So caller MUST set every byte before read.
+func (b *span) Make(_n int) []byte {
+	n := uint32(_n)
+	if n >= b.size || !atomic.CompareAndSwapUint32(&b.lock, 0, 1) {
+		// fallback path: make a new byte slice if current goroutine cannot get the lock or n is out of size
+		return dirtmake.Bytes(int(n), int(n))
+	}
+START:
+	b.read += n
+	// fast path
+	if b.read <= b.size {
+		buf := b.buffer[b.read-n : b.read : b.read]
+		atomic.StoreUint32(&b.lock, 0)
+		return buf
+	}
+	// slow path: create a new buffer
+	b.buffer = dirtmake.Bytes(int(b.size), int(b.size))
+	b.read = 0
+	goto START
+}
+
+// Copy is an alias function for make-and-copy pattern
+func (b *span) Copy(buf []byte) (p []byte) {
+	p = b.Make(len(buf))
+	copy(p, buf)
+	return p
+}
+
+// spanClass calc the minimum number of bits required to represent x
+// [2^sclass,2^(sclass+1)) bytes in a same span class
+func spanClass(size int) int {
+	if size == 0 {
+		return 0
+	}
+	return bits.Len(uint(size))
+}

lang/span/span_test.go

@@ -0,0 +1,161 @@
+// Copyright 2024 ByteDance Inc.
+//
+// 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 span
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSpanClass(t *testing.T) {
+	assert.Equal(t, spanClass(minSpanObject), minSpanClass)
+
+	assert.Equal(t, spanClass(1), 1)
+
+	assert.Equal(t, spanClass(2), 2)
+	assert.Equal(t, spanClass(3), 2)
+
+	assert.Equal(t, spanClass(4), 3)
+	assert.Equal(t, spanClass(5), 3)
+	assert.Equal(t, spanClass(6), 3)
+	assert.Equal(t, spanClass(7), 3)
+
+	assert.Equal(t, spanClass(8), 4)
+	assert.Equal(t, spanClass(9), 4)
+
+	assert.Equal(t, spanClass(32), 6)
+	assert.Equal(t, spanClass(33), 6)
+	assert.Equal(t, spanClass(63), 6)
+
+	assert.Equal(t, spanClass(64), 7)
+	assert.Equal(t, spanClass(65), 7)
+	assert.Equal(t, spanClass(127), 7)
+	assert.Equal(t, spanClass(128), 8)
+	assert.Equal(t, spanClass(129), 8)
+}
+
+func TestSpan(t *testing.T) {
+	bc := NewSpan(1024 * 32)
+
+	var wg sync.WaitGroup
+	for c := 0; c < 12; c++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 1024; i++ {
+				buf := []byte("123")
+				assert.Equal(t, bc.Copy(buf), buf)
+
+				buf = make([]byte, 32)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 63)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 64)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 32*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 64*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 128*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+			}
+		}()
+	}
+	wg.Wait()
+}
+
+func TestSpanCache(t *testing.T) {
+	bc := NewSpanCache(1024 * 32)
+
+	var wg sync.WaitGroup
+	for c := 0; c < 12; c++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 1024; i++ {
+				buf := []byte("123")
+				assert.Equal(t, bc.Copy(buf), buf)
+
+				buf = make([]byte, 32)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 63)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 64)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 32*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 64*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+				buf = make([]byte, 128*1024)
+				assert.Equal(t, len(bc.Copy(buf)), len(buf))
+			}
+		}()
+	}
+	wg.Wait()
+}
+
+var benchStringSizes = []int{
+	15, 31, 127, // not cached
+	128, 256, 1024 - 1, 4*1024 - 1,
+	32*1024 - 1, 128*1024 - 1,
+	512*1024 - 1, // not cached
+}
+
+func BenchmarkSpanCacheCopy(b *testing.B) {
+	for _, sz := range benchStringSizes {
+		b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) {
+			from := make([]byte, sz)
+			sc := NewSpanCache(maxSpanObject * 8)
+			b.RunParallel(func(pb *testing.PB) {
+				b.ReportAllocs()
+				b.ResetTimer()
+				var buffer []byte
+				for pb.Next() {
+					buffer = sc.Copy(from)
+					buffer[0] = 'a'
+					buffer[sz-1] = 'z'
+				}
+				_ = buffer
+			})
+		})
+	}
+}
+
+func BenchmarkMakeAndCopy(b *testing.B) {
+	for _, sz := range benchStringSizes {
+		b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) {
+			from := make([]byte, sz)
+			b.RunParallel(func(pb *testing.PB) {
+				b.ReportAllocs()
+				b.ResetTimer()
+				var buffer []byte
+				for pb.Next() {
+					buffer = make([]byte, sz)
+					copy(buffer, from)
+					buffer[0] = 'a'
+					buffer[sz-1] = 'z'
+				}
+				_ = buffer
+			})
+		})
+	}
+}