storage: reduce intent interleaving allocations

pkg/keys/keys.go

@@ -487,6 +487,24 @@ func LockTableSingleKey(key roachpb.Key, buf []byte) (roachpb.Key, []byte) {
 	return buf, buf
 }
 
+// LockTableSingleNextKey is equivalent to LockTableSingleKey(key.Next(), buf)
+// but avoids an extra allocation in cases where key.Next() must allocate.
+func LockTableSingleNextKey(key roachpb.Key, buf []byte) (roachpb.Key, []byte) {
+	keyLen := len(LocalRangeLockTablePrefix) + len(LockTableSingleKeyInfix) + encoding.EncodeNextBytesSize(key)
+	if cap(buf) < keyLen {
+		buf = make([]byte, 0, keyLen)
+	} else {
+		buf = buf[:0]
+	}
+	// Don't unwrap any local prefix on key using Addr(key). This allow for
+	// doubly-local lock table keys. For example, local range descriptor keys can
+	// be locked during split and merge transactions.
+	buf = append(buf, LocalRangeLockTablePrefix...)
+	buf = append(buf, LockTableSingleKeyInfix...)
+	buf = encoding.EncodeNextBytesAscending(buf, key)
+	return buf, buf
+}
+
 // DecodeLockTableSingleKey decodes the single-key lock table key to return the key
 // that was locked.
 func DecodeLockTableSingleKey(key roachpb.Key) (lockedKey roachpb.Key, err error) {

pkg/keys/keys_test.go

@@ -772,3 +772,28 @@ func TestLockTableKeyEncodeDecode(t *testing.T) {
 		})
 	}
 }
+
+func TestLockTableSingleKeyNext_Equivalent(t *testing.T) {
+	testCases := []struct {
+		key roachpb.Key
+	}{
+		{key: roachpb.Key("foo")},
+		{key: roachpb.Key("a")},
+		{key: roachpb.Key("")},
+		// Causes a doubly-local range local key.
+		{key: RangeDescriptorKey(roachpb.RKey("baz"))},
+	}
+	for _, test := range testCases {
+		t.Run("", func(t *testing.T) {
+			next := test.key.Next()
+			want, _ := LockTableSingleKey(next, nil)
+
+			got, _ := LockTableSingleNextKey(test.key, nil)
+			require.Equal(t, want, got)
+
+			k, err := DecodeLockTableSingleKey(got)
+			require.NoError(t, err)
+			require.Equal(t, next, k)
+		})
+	}
+}

pkg/storage/intent_interleaving_iter.go

@@ -493,7 +493,7 @@ func (i *intentInterleavingIter) SeekGE(key MVCCKey) {
 		intentSeekKey, i.intentKeyBuf = keys.LockTableSingleKey(key.Key, i.intentKeyBuf)
 	} else if !i.prefix {
 		// Seeking to a specific version, so go past the intent.
-		intentSeekKey, i.intentKeyBuf = keys.LockTableSingleKey(key.Key.Next(), i.intentKeyBuf)
+		intentSeekKey, i.intentKeyBuf = keys.LockTableSingleNextKey(key.Key, i.intentKeyBuf)
 	} else {
 		// Else seeking to a particular version and using prefix iteration,
 		// so don't expect to ever see the intent. NB: intentSeekKey is nil.
@@ -1103,7 +1103,7 @@ func (i *intentInterleavingIter) SeekLT(key MVCCKey) {
 		// Seeking to a specific version, so need to see the intent. Since we need
 		// to see the intent for key.Key, and we don't have SeekLE, call Next() on
 		// the key before doing SeekLT.
-		intentSeekKey, i.intentKeyBuf = keys.LockTableSingleKey(key.Key.Next(), i.intentKeyBuf)
+		intentSeekKey, i.intentKeyBuf = keys.LockTableSingleNextKey(key.Key, i.intentKeyBuf)
 	}
 	var limitKey roachpb.Key
 	if i.iterValid {

pkg/util/encoding/encoding.go

@@ -610,6 +610,23 @@ func EncodeBytesAscending(b []byte, data []byte) []byte {
 	return encodeBytesAscendingWithTerminatorAndPrefix(b, data, ascendingBytesEscapes.escapedTerm, bytesMarker)
 }
 
+// EncodeNextBytesAscending encodes the []byte value with an extra 0x00 byte
+// appended before encoding. It's equivalent to
+//
+//	EncodeBytesAscending(b, append(data, 0x00))
+//
+// but may avoid an allocation when the data slice does not have additional
+// capacity.
+func EncodeNextBytesAscending(b []byte, data []byte) []byte {
+	b = append(b, bytesMarker)
+	return encodeNextBytesAscendingWithTerminator(b, data, ascendingBytesEscapes.escapedTerm)
+}
+
+func encodeNextBytesAscendingWithTerminator(b []byte, data []byte, terminator byte) []byte {
+	bs := encodeBytesAscendingWithoutTerminatorOrPrefix(b, data)
+	return append(bs, escape, escaped00, escape, terminator)
+}
+
 // encodeBytesAscendingWithTerminatorAndPrefix encodes the []byte value using an escape-based
 // encoding. The encoded value is terminated with the sequence
 // "\x00\terminator". The encoded bytes are append to the supplied buffer
@@ -676,6 +693,22 @@ func EncodeBytesSize(data []byte) int {
 	return len(data) + 3 + bytes.Count(data, []byte{escape})
 }
 
+// EncodeNextBytesSize returns the size of the []byte value when suffixed with a
+// zero byte and then encoded using EncodeNextBytes{Ascending,Descending}. The
+// function accounts for the encoding marker, escaping, and the terminator.
+func EncodeNextBytesSize(data []byte) int {
+	// Encoding overhead:
+	// +1 for [bytesMarker] prefix
+	// +2 for [escape, escapedTerm] suffix
+	// +1 for each byte that needs to be escaped
+	// +2 for the appended 0x00 byte, plus its escaping byte
+	//
+	// NOTE: bytes.Count is implemented by the go runtime in assembly and is
+	// much faster than looping over the bytes in the slice, especially when
+	// given a single-byte separator.
+	return len(data) + 5 + bytes.Count(data, []byte{escape})
+}
+
 // DecodeBytesAscending decodes a []byte value from the input buffer
 // which was encoded using EncodeBytesAscending. The decoded bytes
 // are appended to r. The remainder of the input buffer and the

pkg/util/encoding/encoding_test.go

@@ -546,6 +546,56 @@ func TestEncodeDecodeBytesAscending(t *testing.T) {
 	}
 }
 
+func TestEncodeNextBytesAscending_Equivalence(t *testing.T) {
+	for _, b := range [][]byte{
+		{0, 1, 'a'},
+		{0, 'a'},
+		{0, 0xff, 'a'},
+		{'a'},
+		{'b'},
+		{'b', 0},
+		{'b', 0, 0},
+		{'b', 0, 0, 'a'},
+		{'b', 0xff},
+		{'h', 'e', 'l', 'l', 'o'},
+	} {
+		next := append(b, 0x00)
+
+		gotSz := EncodeNextBytesSize(b)
+		wantSz := EncodeBytesSize(next)
+		if gotSz != wantSz {
+			t.Errorf("EncodeNextBytesSize(%q) = %d; want %d", b, gotSz, wantSz)
+		}
+		gotV := EncodeNextBytesAscending(nil, b)
+		wantV := EncodeBytesAscending(nil, next)
+		if !bytes.Equal(gotV, wantV) {
+			t.Errorf("EncodeNextBytesAscending(%q) = %q; want %q", b, gotV, wantV)
+		}
+	}
+}
+
+func TestEncodeNextBytesAscending_Equivalence_Randomized(t *testing.T) {
+	rnd, _ := randutil.NewTestRand()
+	var buf [10]byte
+	var nextBuf [10 + 1]byte
+	for i := 0; i < 1000; i++ {
+		b := buf[:randutil.RandIntInRange(rnd, 1, cap(buf))]
+		randutil.ReadTestdataBytes(rnd, b)
+		next := append(append(nextBuf[:0], b...), 0x00)
+
+		gotSz := EncodeNextBytesSize(b)
+		wantSz := EncodeBytesSize(next)
+		if gotSz != wantSz {
+			t.Errorf("EncodeNextBytesSize(%q) = %d; want %d", b, gotSz, wantSz)
+		}
+		gotV := EncodeNextBytesAscending(nil, b)
+		wantV := EncodeBytesAscending(nil, next)
+		if !bytes.Equal(gotV, wantV) {
+			t.Errorf("EncodeNextBytesAscending(%q) = %q; want %q", b, gotV, wantV)
+		}
+	}
+}
+
 func TestEncodeDecodeBytesDescending(t *testing.T) {
 	testCases := []struct {
 		value   []byte