diff --git a/cralloc/scratch_buffer.go b/cralloc/scratch_buffer.go new file mode 100644 index 0000000..07a3bee --- /dev/null +++ b/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{} + } +} diff --git a/cralloc/scratch_buffer_test.go b/cralloc/scratch_buffer_test.go new file mode 100644 index 0000000..484387b --- /dev/null +++ b/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) + }) +}