diff --git a/pkg/keys/keys.go b/pkg/keys/keys.go index f2c2dcb26351..ee2b4564fae0 100644 --- a/pkg/keys/keys.go +++ b/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) { diff --git a/pkg/keys/keys_test.go b/pkg/keys/keys_test.go index 39957add50c5..f63c3315197b 100644 --- a/pkg/keys/keys_test.go +++ b/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) + }) + } +} diff --git a/pkg/storage/intent_interleaving_iter.go b/pkg/storage/intent_interleaving_iter.go index 80be79c9c1ab..e3a333e46a81 100644 --- a/pkg/storage/intent_interleaving_iter.go +++ b/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 { diff --git a/pkg/util/encoding/encoding.go b/pkg/util/encoding/encoding.go index b95029eb87e9..a9716372fda1 100644 --- a/pkg/util/encoding/encoding.go +++ b/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 diff --git a/pkg/util/encoding/encoding_test.go b/pkg/util/encoding/encoding_test.go index 0bf14d9a503f..d29d5b5829b3 100644 --- a/pkg/util/encoding/encoding_test.go +++ b/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