sstable: add (*sstable.Writer).RangeKey{Set,Unset,Delete} methods

Users of `sstable.Writer` other than Pebble itself (e.g. Cockroach)
require a means of adding range keys to an sstable that is slightly more
flexible than the existing `AddRangeKey` method in the way in which key
spans can be added to the table.

Specifically, `AddRangeKey` requires that spans be sorted (start / end
key ascending, sequence number and key kind descending), in addition to
fragmented and coalesced (i.e. multiple suffixes within the same span
are coalesced into a single value).

Add the `RangeKey{Set,Unset,Delete}` methods on the `sstable.Writer`
that allow adding spans in order of start key. These methods have no
requirement that the spans be fragmented or coalesced. Instead, the
implementation handles the fragmenting of spans as keys are added. A
`keyspan.Fragmenter` emits completed spans into a `rangekey.coalescer`,
which aggregates the fragments for the same span into a
`rangekey.CoalescedSpan`, which can then be used to write the individual
`RangeKey{Set,Unset,Delete}` key / value pairs into the rang key block
in the sstable.

The new `RangeKey{Set,Unset,Delete}` methods have similar semantics to
the existing `Set,Delete,DeleteRange,Merge` methods that write the key /
value pairs at sequence number zero. In the range key case, these new
methods are intended to be used by external callers when preparing
sstables for ingestion.

Add a data-driven test specific to the new range key methods. These
tests are separate to the existing data driven test cases for the
`AddRangeKey` method, which has stricter requirements on how keys are
added (i.e. already ordered, fragmented and coalesced).

Rename `AddInternalRangeKey` to `AddRangeKey` to better match the
existing naming of `(*sstable.Writer).Add`.

Add new utility methods to the `rangekey` package to aid in encoding
range key values.

Related to #1339.

internal/rangekey/rangekey.go

@@ -67,32 +67,28 @@ type SuffixValue struct {
 	Value  []byte
 }
 
-// EncodedSetValueLen precomputes the length of a RangeKeySet's value when
-// encoded.  It may be used to construct a buffer of the appropriate size before
-// encoding.
-func EncodedSetValueLen(endKey []byte, suffixValues []SuffixValue) int {
-	n := lenVarint(len(endKey))
-	n += len(endKey)
+// EncodedSetSuffixValuesLen precomputes the length of the given slice of
+// SuffixValues, when encoded for a RangeKeySet. It may be used to construct a
+// buffer of the appropriate size before encoding.
+func EncodedSetSuffixValuesLen(suffixValues []SuffixValue) int {
+	var n int
 	for i := 0; i < len(suffixValues); i++ {
 		n += lenVarint(len(suffixValues[i].Suffix))
 		n += len(suffixValues[i].Suffix)
 		n += lenVarint(len(suffixValues[i].Value))
 		n += len(suffixValues[i].Value)
 	}
-
 	return n
 }
 
-// EncodeSetValue encodes a RangeKeySet's value into dst. The length of dst must
-// be greater than or equal to EncodedSetValueLen. EncodeSetValue returns the
-// number of bytes written, which should always equal the EncodedSetValueLen
-// with the same arguments.
-func EncodeSetValue(dst []byte, endKey []byte, suffixValues []SuffixValue) int {
-	// First encode the end key as a varstring.
-	n := binary.PutUvarint(dst, uint64(len(endKey)))
-	n += copy(dst[n:], endKey)
-
+// EncodeSetSuffixValues encodes a slice of SuffixValues for a RangeKeySet into
+// dst. The length of dst must be greater than or equal to
+// EncodedSetSuffixValuesLen. EncodeSetSuffixValues returns the number of bytes
+// written, which should always equal the EncodedSetValueLen with the same
+// arguments.
+func EncodeSetSuffixValues(dst []byte, suffixValues []SuffixValue) int {
 	// Encode the list of (suffix, value-len) tuples.
+	var n int
 	for i := 0; i < len(suffixValues); i++ {
 		// Encode the length of the suffix.
 		n += binary.PutUvarint(dst[n:], uint64(len(suffixValues[i].Suffix)))
@@ -109,6 +105,28 @@ func EncodeSetValue(dst []byte, endKey []byte, suffixValues []SuffixValue) int {
 	return n
 }
 
+// EncodedSetValueLen precomputes the length of a RangeKeySet's value when
+// encoded. It may be used to construct a buffer of the appropriate size before
+// encoding.
+func EncodedSetValueLen(endKey []byte, suffixValues []SuffixValue) int {
+	n := lenVarint(len(endKey))
+	n += len(endKey)
+	n += EncodedSetSuffixValuesLen(suffixValues)
+	return n
+}
+
+// EncodeSetValue encodes a RangeKeySet's value into dst. The length of dst must
+// be greater than or equal to EncodedSetValueLen. EncodeSetValue returns the
+// number of bytes written, which should always equal the EncodedSetValueLen
+// with the same arguments.
+func EncodeSetValue(dst []byte, endKey []byte, suffixValues []SuffixValue) int {
+	// First encode the end key as a varstring.
+	n := binary.PutUvarint(dst, uint64(len(endKey)))
+	n += copy(dst[n:], endKey)
+	n += EncodeSetSuffixValues(dst[n:], suffixValues)
+	return n
+}
+
 // DecodeEndKey reads the end key from the beginning of a range key (RANGEKEYSET,
 // RANGEKEYUNSET or RANGEKEYDEL)'s physical encoded value. Both sets and unsets
 // encode the range key, plus additional data in the value.
@@ -147,17 +165,42 @@ func DecodeSuffixValue(data []byte) (sv SuffixValue, rest []byte, ok bool) {
 	return sv, data, true
 }
 
+// EncodedUnsetSuffixesLen precomputes the length of the given slice of
+// suffixes, when encoded for a RangeKeyUnset. It may be used to construct a
+// buffer of the appropriate size before encoding.
+func EncodedUnsetSuffixesLen(suffixes [][]byte) int {
+	var n int
+	for i := 0; i < len(suffixes); i++ {
+		n += lenVarint(len(suffixes[i]))
+		n += len(suffixes[i])
+	}
+	return n
+}
+
+// EncodeUnsetSuffixes encodes a slice of suffixes for a RangeKeyUnset into dst.
+// The length of dst must be greater than or equal to EncodedUnsetSuffixesLen.
+// EncodeUnsetSuffixes returns the number of bytes written, which should always
+// equal the EncodedUnsetSuffixesLen with the same arguments.
+func EncodeUnsetSuffixes(dst []byte, suffixes [][]byte) int {
+	// Encode the list of (suffix, value-len) tuples.
+	var n int
+	for i := 0; i < len(suffixes); i++ {
+		//  Encode the length of the suffix.
+		n += binary.PutUvarint(dst[n:], uint64(len(suffixes[i])))
+
+		// Encode the suffix itself.
+		n += copy(dst[n:], suffixes[i])
+	}
+	return n
+}
+
 // EncodedUnsetValueLen precomputes the length of a RangeKeyUnset's value when
 // encoded.  It may be used to construct a buffer of the appropriate size before
 // encoding.
 func EncodedUnsetValueLen(endKey []byte, suffixes [][]byte) int {
 	n := lenVarint(len(endKey))
 	n += len(endKey)
-
-	for i := 0; i < len(suffixes); i++ {
-		n += lenVarint(len(suffixes[i]))
-		n += len(suffixes[i])
-	}
+	n += EncodedUnsetSuffixesLen(suffixes)
 	return n
 }
 
@@ -169,15 +212,7 @@ func EncodeUnsetValue(dst []byte, endKey []byte, suffixes [][]byte) int {
 	// First encode the end key as a varstring.
 	n := binary.PutUvarint(dst, uint64(len(endKey)))
 	n += copy(dst[n:], endKey)
-
-	// Encode the list of suffix varstrings.
-	for i := 0; i < len(suffixes); i++ {
-		//  Encode the length of the suffix.
-		n += binary.PutUvarint(dst[n:], uint64(len(suffixes[i])))
-
-		// Encode the suffix itself.
-		n += copy(dst[n:], suffixes[i])
-	}
+	n += EncodeUnsetSuffixes(dst[n:], suffixes)
 	return n
 }
 

internal/rangekey/rangekey_test.go

@@ -11,6 +11,47 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestSetSuffixValues_RoundTrip(t *testing.T) {
+	testCases := [][]SuffixValue{
+		{
+			{Suffix: []byte{}, Value: []byte("world")},
+		},
+		{
+			{Suffix: []byte("foo"), Value: []byte("bar")},
+		},
+		{
+			{Suffix: []byte(""), Value: []byte("boop")},
+			{Suffix: []byte("foo"), Value: []byte("beep")},
+			{Suffix: []byte("bar"), Value: []byte("bop")},
+			{Suffix: []byte("bax"), Value: []byte("boink")},
+			{Suffix: []byte("zoop"), Value: []byte("zoink")},
+		},
+	}
+
+	var b []byte
+	for _, tc := range testCases {
+		// Encode.
+		l := EncodedSetSuffixValuesLen(tc)
+		if l <= cap(b) {
+			b = b[:l]
+		} else {
+			b = make([]byte, l)
+		}
+		n := EncodeSetSuffixValues(b, tc)
+		require.Equal(t, l, n)
+
+		// Decode.
+		var suffixValues []SuffixValue
+		for len(b) > 0 {
+			sv, rest, ok := DecodeSuffixValue(b)
+			require.True(t, ok)
+			suffixValues = append(suffixValues, sv)
+			b = rest
+		}
+		require.Equal(t, tc, suffixValues)
+	}
+}
+
 func TestSetValue_Roundtrip(t *testing.T) {
 	testCases := []struct {
 		endKey       []byte
@@ -71,6 +112,44 @@ func TestSetValue_Roundtrip(t *testing.T) {
 	}
 }
 
+func TestUnsetSuffixes_RoundTrip(t *testing.T) {
+	type suffixes [][]byte
+	testCases := []suffixes{
+		{{}},
+		{[]byte("foo")},
+		{
+			{},
+			[]byte("foo"),
+			[]byte("bar"),
+			[]byte("bax"),
+			[]byte("zoop"),
+		},
+	}
+
+	var b []byte
+	for _, tc := range testCases {
+		// Encode.
+		l := EncodedUnsetSuffixesLen(tc)
+		if l <= cap(b) {
+			b = b[:l]
+		} else {
+			b = make([]byte, l)
+		}
+		n := EncodeUnsetSuffixes(b, tc)
+		require.Equal(t, l, n)
+
+		// Decode.
+		var ss suffixes
+		for len(b) > 0 {
+			s, rest, ok := DecodeSuffix(b)
+			require.True(t, ok)
+			ss = append(ss, s)
+			b = rest
+		}
+		require.Equal(t, tc, ss)
+	}
+}
+
 func TestUnsetValue_Roundtrip(t *testing.T) {
 	testCases := []struct {
 		endKey   []byte

sstable/data_test.go

@@ -142,7 +142,7 @@ func runBuildCmd(
 		n := rangekey.RecombinedValueLen(v.Start.Kind(), v.End, v.Value)
 		b := make([]byte, n)
 		_ = rangekey.RecombineValue(v.Start.Kind(), b, v.End, v.Value)
-		if err := w.AddInternalRangeKey(v.Start, b); err != nil {
+		if err := w.AddRangeKey(v.Start, b); err != nil {
 			return nil, nil, err
 		}
 	}
@@ -198,7 +198,7 @@ func runBuildRawCmd(td *datadriven.TestData) (*WriterMetadata, *Reader, error) {
 			// Values for range keys must be converted into their "packed" form before
 			// being added to the Writer.
 			_, value := rangekey.Parse(data)
-			if err := w.AddInternalRangeKey(key, value); err != nil {
+			if err := w.AddRangeKey(key, value); err != nil {
 				return nil, nil, err
 			}
 		default: