cralloc: add ScratchBuffer

cralloc/scratch_buffer.go

@@ -0,0 +1,110 @@
+// Copyright 2025 The Cockroach 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 cralloc
+
+import "unsafe"
+
+// ScratchBuffer is a helper for the common pattern of reusing a byte buffer to
+// reduce slice allocations. To use, replace `make([]byte, n)` with
+// `sb.Alloc(n)`.
+type ScratchBuffer struct {
+	p        unsafe.Pointer
+	capacity int
+}
+
+// AllocUnsafe returns a byte slice of length n and arbitrary capacity which can
+// be used until the next call to AllocUnsafe/AllocZero/Append.
+//
+// WARNING: the slice contains arbitrary data.
+//
+// This method is marked Unsafe because the allowed lifetime of the returned
+// slice is limited.
+//
+// If the receiver is nil, always allocates a new slice.
+func (sb *ScratchBuffer) AllocUnsafe(n int) []byte {
+	if sb == nil {
+		return make([]byte, n)
+	}
+	s := unsafe.Slice((*byte)(sb.p), sb.capacity)
+	if sb.capacity >= n {
+		return s[:n]
+	}
+	// Adapted from slices.Grow().
+	s = append(s[:0], make([]byte, n)...)
+	sb.p = unsafe.Pointer(&s[0])
+	sb.capacity = cap(s)
+	return s
+}
+
+// AllocZeroUnsafe returns a byte slice of length n and arbitrary capacity which
+// can be used until the next call to AllocUnsafe/AllocZero/Append. The slice is
+// zeroed out.
+//
+// WARNING: the slice contains arbitrary data between the length and the
+// capacity.
+//
+// This method is marked Unsafe because the allowed lifetime of the returned
+// slice is limited.
+//
+// If the receiver is nil, always allocates a new slice.
+func (sb *ScratchBuffer) AllocZeroUnsafe(n int) []byte {
+	if sb == nil {
+		return make([]byte, n)
+	}
+	s := unsafe.Slice((*byte)(sb.p), sb.capacity)
+	if sb.capacity >= n {
+		s = s[:n]
+		clear(s)
+		return s
+	}
+	// Adapted from slices.Grow(). We do not want to simply use make([]byte, n)
+	// because we want the scratch buffer to grow according to the append()
+	// heuristics. Otherwise, an allocation pattern of slowly increasing sizes
+	// would cause an allocation each time.
+	s = append(s[:0], make([]byte, n)...)
+	sb.p = unsafe.Pointer(&s[0])
+	sb.capacity = cap(s)
+	return s
+}
+
+// Append is like the built-in append(), but it also updates the scratch buffer
+// so that any newly allocated buffer can be reused.
+//
+// Append can be used with buffers not allocated through the scratch buffer (in
+// which case the scratch buffer is not updated).
+func (sb *ScratchBuffer) Append(buf []byte, values ...byte) []byte {
+	res := append(buf, values...)
+	if sb != nil && unsafe.SliceData(buf) == (*byte)(sb.p) && unsafe.SliceData(res) != (*byte)(sb.p) {
+		sb.p = unsafe.Pointer(unsafe.SliceData(res))
+		sb.capacity = cap(res)
+	}
+	return res
+}
+
+// Capacity returns the current capacity.
+func (sb *ScratchBuffer) Capacity() int {
+	if sb == nil {
+		return 0
+	}
+	return sb.capacity
+}
+
+// Reset clears the buffer. This can be useful if we want to avoid retaining a
+// very large buffer.
+func (sb *ScratchBuffer) Reset() {
+	if sb != nil {
+		*sb = ScratchBuffer{}
+	}
+}

cralloc/scratch_buffer_test.go

@@ -0,0 +1,71 @@
+// Copyright 2025 The Cockroach 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 cralloc
+
+import (
+	"math/rand/v2"
+	"testing"
+
+	"github.com/cockroachdb/crlib/testutils/require"
+)
+
+func TestScratchBuffer(t *testing.T) {
+	s := (*ScratchBuffer)(nil).AllocUnsafe(100)
+	require.Equal(t, len(s), 100)
+	var sb ScratchBuffer
+	s = sb.AllocUnsafe(100)
+	require.Equal(t, len(s), 100)
+	c := cap(s)
+	s = sb.AllocUnsafe(50)
+	require.Equal(t, len(s), 50)
+	require.Equal(t, cap(s), c)
+	s = sb.AllocUnsafe(101)
+	require.Equal(t, len(s), 101)
+	require.GT(t, cap(s), 101)
+
+	t.Run("AllocZero", func(t *testing.T) {
+		for range 100 {
+			var sb ScratchBuffer
+			maxN := 1 + rand.IntN(1000)
+			for range 20 {
+				n := rand.IntN(maxN)
+				b := sb.AllocZeroUnsafe(n)
+				for i := range b {
+					require.Equal(t, b[i], 0)
+				}
+				// Trash the entire buffer.
+				b = b[:cap(b)]
+				for i := range b {
+					b[i] = 0xcc
+				}
+			}
+		}
+	})
+
+	t.Run("Append", func(t *testing.T) {
+		var sb ScratchBuffer
+		b := sb.AllocUnsafe(100)
+		b = sb.Append(b, make([]byte, 1000)...)
+		require.Equal(t, len(b), 1100)
+		// Ensure the capacity has grown.
+		require.GE(t, sb.Capacity(), 1100)
+
+		// Append an unrelated slice.
+		b = sb.Append(make([]byte, 1100), make([]byte, 10000)...)
+		require.Equal(t, len(b), 11100)
+		// Ensure the capacity did not grow.
+		require.LT(t, sb.Capacity(), 10000)
+	})
+}