sstable/block: fix Buffer growth bug

Fix a bug whereby Buffer.AppendValue would fail to sufficiently grow the
capacity of its backing byte slice. Previously, AppendValue used cap(buf) when
calculating the argument to slices.Grow, but slices.Grow expects its argument
to be in terms of the slice's current length, not capacity.

Add a focused randomized test that catches the bug. Also, add invariant
assertions around a BufHandle's existence in or out of a pool.

Fix #4247.

Author: jbowens

sstable/block/compression.go

@@ -338,7 +338,7 @@ func (b *Buffer) Append(v []byte) int {
 		for size < newLen {
 			size *= 2
 		}
-		b.h.b = slices.Grow(b.h.b, size-cap(b.h.b))
+		b.h.b = slices.Grow(b.h.b, size-len(b.h.b))
 	}
 	b.h.b = b.h.b[:newLen]
 	if n := copy(b.h.b[off:], v); n != len(v) {
@@ -411,10 +411,10 @@ func (h *BufHandle) Release() {
 		panic(errors.AssertionFailedf("pool has no maximum size"))
 	}
 	// Note we avoid releasing buffers that are larger than the configured
-	// maximum to the pool. This avoids holding on to occassional large buffers
-	// necesary for, for example, single large values.
+	// maximum to the pool. This avoids holding on to occasional large buffers
+	// necessary for, for example, singlular large values.
 	if h.b != nil && len(h.b) < h.pool.Max {
-		if invariants.Enabled {
+		if invariants.Sometimes(50) {
 			// Set the bytes to a random value. Cap the number of bytes being
 			// randomized to prevent test timeouts.
 			l := min(cap(h.b), 1000)
@@ -423,7 +423,7 @@ func (h *BufHandle) Release() {
 				h.b[j] = byte(rand.Uint32())
 			}
 		}
-		h.pool.pool.Put(h)
+		h.pool.Put(h)
 	}
 }
 
@@ -446,12 +446,29 @@ type bufferSyncPool struct {
 	pool    sync.Pool
 }
 
+// Put returns a buffer to the pool. While the buffer is in the pool, its pool
+// member is zeroed. This is used to validate invariants around double use of a
+// buffer.
+func (p *bufferSyncPool) Put(bh *BufHandle) {
+	if bh.pool != p {
+		panic(errors.AssertionFailedf("buffer has pool %v; trying to return it to pool %v", bh.pool, p))
+	}
+	bh.pool = nil
+	p.pool.Put(bh)
+}
+
 // Get retrieves a new buf from the pool, or allocates one of the configured
 // default size if the pool is empty.
 func (p *bufferSyncPool) Get() *BufHandle {
 	v := p.pool.Get()
 	if v != nil {
-		return v.(*BufHandle)
+		bh := v.(*BufHandle)
+		if bh.pool != nil {
+			panic(errors.AssertionFailedf("buffer has a pool; was it inserted into a pool twice?"))
+		}
+		// Set the pool so we know where to return the buffer to.
+		bh.pool = p
+		return bh
 	}
 	if invariants.Enabled && p.Default == 0 {
 		// Guard against accidentally forgetting to initialize a buffer sync pool.

sstable/block/compression_test.go

@@ -6,6 +6,7 @@ package block
 
 import (
 	"encoding/binary"
+	"fmt"
 	"math/rand/v2"
 	"testing"
 	"time"
@@ -92,3 +93,45 @@ func decompress(algo CompressionIndicator, b []byte) (*cache.Value, error) {
 	}
 	return decoded, nil
 }
+
+func TestBufferRandomized(t *testing.T) {
+	seed := uint64(time.Now().UnixNano())
+	t.Logf("seed %d", seed)
+	rng := rand.New(rand.NewPCG(0, seed))
+
+	var b Buffer
+	b.Init(SnappyCompression, ChecksumTypeCRC32c)
+	defer b.Release()
+	vbuf := make([]byte, 0, 1<<10) // 1 KiB
+
+	for i := 0; i < 25; i++ {
+		t.Run(fmt.Sprintf("iteration %d", i), func(t *testing.T) {
+			// Randomly release and reinitialize the buffer.
+			if rng.IntN(5) == 1 {
+				b.Release()
+				b.Init(SnappyCompression, ChecksumTypeCRC32c)
+			}
+
+			aggregateSizeOfKVs := rng.IntN(4<<20-(1<<10)) + 1<<10 // [1 KiB, 4 MiB)
+			size := 0
+			for b.Size() < aggregateSizeOfKVs {
+				vlen := rng.IntN(aggregateSizeOfKVs-b.Size()) + 1
+				if cap(vbuf) < vlen {
+					vbuf = make([]byte, vlen)
+				} else {
+					vbuf = vbuf[:vlen]
+				}
+				for i := range vbuf {
+					vbuf[i] = byte(rng.Uint32())
+				}
+				b.Append(vbuf)
+				size += vlen
+				require.Equal(t, size, b.Size())
+				s := b.Get()
+				require.Equal(t, vbuf, s[len(s)-len(vbuf):])
+			}
+			_, bh := b.CompressAndChecksum()
+			bh.Release()
+		})
+	}
+}