colblk: use regular slice in PrefixBytesIter

We currently store an `unsafe.Pointer` and separate length and
capacities. In many cases in the hot path we have to create a slice
using `unsafe.Slice`. However `unsafe.Slice` includes some checks that
we don't need (see [here](https://github.com/golang/go/blob/f38d42f2c4c6ad0d7cbdad5e1417cac3be2a5dcb/src/runtime/unsafe.go#L53)).

In this change we switch to storing a regular slice and converting to
an `unsafe.Pointer` using `unsafe.SliceData` (which directly accesses
the pointer field). This improves performance, even while gaining
boundary checks when appending the suffix.

```
CockroachDataBlockIter/AlphaLen=4,Shared=8,PrefixLen=32,Logical=0,value=8/Next    11.4ns ± 2%    10.6ns ± 1%  -7.32%  (p=0.008 n=5+5)
```

sstable/colblk/cockroach_test.go

@@ -270,25 +270,42 @@ func (ks *cockroachKeySeeker) MaterializeUserKey(ki *PrefixBytesIter, prevRow, r
 		ks.prefixes.SetAt(ki, row)
 	}
 
+	ptr := unsafe.Pointer(uintptr(unsafe.Pointer(unsafe.SliceData(ki.buf))) + uintptr(len(ki.buf)))
 	mvccWall := ks.mvccWallTimes.At(row)
 	mvccLogical := uint32(ks.mvccLogical.At(row))
 	if mvccWall == 0 && mvccLogical == 0 {
 		// This is not an MVCC key. Use the untyped suffix.
 		untypedSuffixed := ks.untypedSuffixes.At(row)
-		memmove(unsafe.Pointer(uintptr(ki.ptr)+uintptr(ki.len)), unsafe.Pointer(unsafe.SliceData(untypedSuffixed)), uintptr(len(untypedSuffixed)))
-		return unsafe.Slice((*byte)(ki.ptr), ki.len+len(untypedSuffixed))
+		res := ki.buf[:len(ki.buf)+len(untypedSuffixed)]
+		memmove(ptr, unsafe.Pointer(unsafe.SliceData(untypedSuffixed)), uintptr(len(untypedSuffixed)))
+		return res
 	}
 
+	// Inline binary.BigEndian.PutUint64. Note that this code is converted into
+	// word-size instructions by the compiler.
+	*(*byte)(ptr) = byte(mvccWall >> 56)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 1)) = byte(mvccWall >> 48)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 2)) = byte(mvccWall >> 40)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 3)) = byte(mvccWall >> 32)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 4)) = byte(mvccWall >> 24)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 5)) = byte(mvccWall >> 16)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 6)) = byte(mvccWall >> 8)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 7)) = byte(mvccWall)
+
+	ptr = unsafe.Pointer(uintptr(ptr) + 8)
 	// This is an MVCC key.
 	if mvccLogical == 0 {
-		binary.BigEndian.PutUint64(unsafe.Slice((*byte)(unsafe.Pointer(uintptr(ki.ptr)+uintptr(ki.len))), 8), mvccWall)
-		*(*byte)(unsafe.Pointer(uintptr(ki.ptr) + uintptr(ki.len) + 8)) = 9
-		return unsafe.Slice((*byte)(ki.ptr), ki.len+9)
+		*(*byte)(ptr) = 9
+		return ki.buf[:len(ki.buf)+9]
 	}
-	binary.BigEndian.PutUint64(unsafe.Slice((*byte)(unsafe.Pointer(uintptr(ki.ptr)+uintptr(ki.len))), 8), mvccWall)
-	binary.BigEndian.PutUint32(unsafe.Slice((*byte)(unsafe.Pointer(uintptr(ki.ptr)+uintptr(ki.len+8))), 4), mvccLogical)
-	*(*byte)(unsafe.Pointer(uintptr(ki.ptr) + uintptr(ki.len) + 12)) = 13
-	return unsafe.Slice((*byte)(ki.ptr), ki.len+13)
+
+	// Inline binary.BigEndian.PutUint32.
+	*(*byte)(ptr) = byte(mvccWall >> 24)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 1)) = byte(mvccWall >> 16)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 2)) = byte(mvccWall >> 8)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 3)) = byte(mvccWall)
+	*(*byte)(unsafe.Pointer(uintptr(ptr) + 4)) = 13
+	return ki.buf[:len(ki.buf)+13]
 }
 
 // Release is part of the KeySeeker interface.

sstable/colblk/data_block.go

@@ -327,8 +327,13 @@ func (ks *defaultKeySeeker) MaterializeUserKey(keyIter *PrefixBytesIter, prevRow
 		ks.prefixes.SetAt(keyIter, row)
 	}
 	suffix := ks.suffixes.At(row)
-	memmove(unsafe.Pointer(uintptr(keyIter.ptr)+uintptr(keyIter.len)), unsafe.Pointer(unsafe.SliceData(suffix)), uintptr(len(suffix)))
-	return unsafe.Slice((*byte)(keyIter.ptr), keyIter.len+len(suffix))
+	res := keyIter.buf[:len(keyIter.buf)+len(suffix)]
+	memmove(
+		unsafe.Pointer(uintptr(unsafe.Pointer(unsafe.SliceData(keyIter.buf)))+uintptr(len(keyIter.buf))),
+		unsafe.Pointer(unsafe.SliceData(suffix)),
+		uintptr(len(suffix)),
+	)
+	return res
 }
 
 func (ks *defaultKeySeeker) Release() {
@@ -596,7 +601,10 @@ func (i *DataBlockIter) Init(
 	}
 	// Allocate a keyIter buffer that's large enough to hold the largest user
 	// key in the block.
-	i.keyIter.Alloc(int(r.maximumKeyLength))
+
+	n := int(r.maximumKeyLength)
+	ptr := mallocgc(uintptr(n), nil, false)
+	i.keyIter.buf = unsafe.Slice((*byte)(ptr), n)[:0]
 	return i.keySeeker.Init(r)
 }
 

sstable/colblk/prefix_bytes.go

@@ -239,8 +239,11 @@ func (b PrefixBytes) At(i int) []byte {
 // within a PrefixBytes, avoiding unnecessary copying when portions of slices
 // are shared.
 type PrefixBytesIter struct {
-	UnsafeBuf
-	bundlePrefixLen       uint32
+	// buf is used for materializing a user key. It is preallocated to the maximum
+	// key length in the data block.
+	buf             []byte
+	bundlePrefixLen uint32
+
 	offsetIndex           int
 	nextBundleOffsetIndex int
 }
@@ -264,21 +267,19 @@ func (b *PrefixBytes) SetAt(it *PrefixBytesIter, i int) {
 	rowSuffixStart, rowSuffixEnd := b.rowSuffixOffsets(i, it.offsetIndex)
 
 	// Grow the size of the iterator's buffer if necessary.
-	it.len = b.sharedPrefixLen + int(it.bundlePrefixLen) + int(rowSuffixEnd-rowSuffixStart)
-	if it.len > it.cap {
-		panic(errors.AssertionFailedf("buffer too small: %d > %d", it.len, it.cap))
-	}
+	it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)]
 
+	ptr := unsafe.Pointer(unsafe.SliceData(it.buf))
 	// Copy the shared key prefix.
-	memmove(it.ptr, b.rawBytes.data, uintptr(b.sharedPrefixLen))
+	memmove(ptr, b.rawBytes.data, uintptr(b.sharedPrefixLen))
 	// Copy the bundle prefix.
 	memmove(
-		unsafe.Pointer(uintptr(it.ptr)+uintptr(b.sharedPrefixLen)),
+		unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)),
 		unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(bundleOffsetStart)),
 		uintptr(it.bundlePrefixLen))
 	// Copy the per-row suffix.
 	memmove(
-		unsafe.Pointer(uintptr(it.ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
+		unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
 		unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)),
 		uintptr(rowSuffixEnd-rowSuffixStart))
 	// Set nextBundleOffsetIndex so that a call to SetNext can cheaply determine
@@ -305,13 +306,11 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) {
 			return
 		}
 		// Grow the buffer if necessary.
-		it.len = b.sharedPrefixLen + int(it.bundlePrefixLen) + int(rowSuffixEnd-rowSuffixStart)
-		if it.len > it.cap {
-			panic(errors.AssertionFailedf("buffer too small: %d > %d", it.len, it.cap))
-		}
+		it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)]
 		// Copy in the per-row suffix.
+		ptr := unsafe.Pointer(unsafe.SliceData(it.buf))
 		memmove(
-			unsafe.Pointer(uintptr(it.ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
+			unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
 			unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)),
 			uintptr(rowSuffixEnd-rowSuffixStart))
 		return
@@ -331,18 +330,16 @@ func (b *PrefixBytes) SetNext(it *PrefixBytesIter) {
 	it.nextBundleOffsetIndex = it.offsetIndex + (1 << b.bundleShift)
 
 	// Grow the buffer if necessary.
-	it.len = b.sharedPrefixLen + int(it.bundlePrefixLen) + int(rowSuffixEnd-rowSuffixStart)
-	if it.len > it.cap {
-		panic(errors.AssertionFailedf("buffer too small: %d > %d", it.len, it.cap))
-	}
+	it.buf = it.buf[:b.sharedPrefixLen+int(it.bundlePrefixLen)+int(rowSuffixEnd-rowSuffixStart)]
 	// Copy in the new bundle suffix.
+	ptr := unsafe.Pointer(unsafe.SliceData(it.buf))
 	memmove(
-		unsafe.Pointer(uintptr(it.ptr)+uintptr(b.sharedPrefixLen)),
+		unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)),
 		unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(bundlePrefixStart)),
 		uintptr(it.bundlePrefixLen))
 	// Copy in the per-row suffix.
 	memmove(
-		unsafe.Pointer(uintptr(it.ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
+		unsafe.Pointer(uintptr(ptr)+uintptr(b.sharedPrefixLen)+uintptr(it.bundlePrefixLen)),
 		unsafe.Pointer(uintptr(b.rawBytes.data)+uintptr(rowSuffixStart)),
 		uintptr(rowSuffixEnd-rowSuffixStart))
 }

sstable/colblk/prefix_bytes_test.go

@@ -276,17 +276,17 @@ func BenchmarkPrefixBytes(b *testing.B) {
 			pb, _ := DecodePrefixBytes(buf, 0, n)
 			b.ResetTimer()
 			var pbi PrefixBytesIter
-			pbi.Alloc(maxLen)
+			pbi.buf = make([]byte, 0, maxLen)
 			for i := 0; i < b.N; i++ {
 				j := i % n
 				if j == 0 {
 					pb.SetAt(&pbi, j)
 				} else {
 					pb.SetNext(&pbi)
 				}
-				if invariants.Enabled && !bytes.Equal(pbi.UnsafeSlice(), userKeys[j]) {
+				if invariants.Enabled && !bytes.Equal(pbi.buf, userKeys[j]) {
 					b.Fatalf("Constructed key %q (%q, %q, %q) for index %d; expected %q",
-						pbi.UnsafeSlice(), pb.SharedPrefix(), pb.RowBundlePrefix(j), pb.RowSuffix(j), j, userKeys[j])
+						pbi.buf, pb.SharedPrefix(), pb.RowBundlePrefix(j), pb.RowSuffix(j), j, userKeys[j])
 				}
 			}
 		})