Repository: cockroachdb/pebble
Pull request #3840 (2024-08-12 18:05:18)

db: various DB.SSTables improvements
Add a couple improvements to the DB.SSTables method that powers the crdb_internal.sstable_metrics builtin in Cockroach. Subsequent work up in Cockroach will adopt a custom JSON serialization to improve the ergonomics of working with the builtin

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/3840)
<!-- Reviewable:end -->

jbowens: Removed my datadriven test because it's currently nondeterministic. I'll revive it in the Cockroach repo so that I can test the formatting that will actually be in use by the sstable_metrics builtin.

Patch:
diff --git a/db.go b/db.go
index 2760f2cc..8c7b6f16 100644
--- a/db.go
+++ b/db.go
@@ -9,7 +9,6 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"strconv"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -2072,8 +2071,7 @@ func WithKeyRangeFilter(start, end []byte) SSTablesOption {
 
 // WithApproximateSpanBytes enables capturing the approximate number of bytes that
 // overlap the provided key span for each sstable.
-// NOTE: this option can only be used with WithKeyRangeFilter and WithProperties
-// provided.
+// NOTE: This option requires WithKeyRangeFilter.
 func WithApproximateSpanBytes() SSTablesOption {
 	return func(opt *sstablesOptions) {
 		opt.withApproximateSpanBytes = true
@@ -2100,8 +2098,21 @@ const (
 	// or lifecycle management. An example of an external file is a file restored
 	// from a backup.
 	BackingTypeExternal
+	backingTypeCount
 )
 
+var backingTypeToString = [backingTypeCount]string{
+	BackingTypeLocal:         "local",
+	BackingTypeShared:        "shared",
+	BackingTypeSharedForeign: "shared-foreign",
+	BackingTypeExternal:      "external",
+}
+
+// String implements fmt.Stringer.
+func (b BackingType) String() string {
+	return backingTypeToString[b]
+}
+
 // SSTableInfo export manifest.TableInfo with sstable.Properties alongside
 // other file backing info.
 type SSTableInfo struct {
@@ -2116,6 +2127,10 @@ type SSTableInfo struct {
 	// Locator is the remote.Locator backing this sstable, if the backing type is
 	// not BackingTypeLocal.
 	Locator remote.Locator
+	// ApproximateSpanBytes describes the approximate number of bytes within the
+	// sstable that fall within a particular span. It's populated only when the
+	// ApproximateSpanBytes option is passed into DB.SSTables.
+	ApproximateSpanBytes uint64 `json:"ApproximateSpanBytes,omitempty"`
 
 	// Properties is the sstable properties of this table. If Virtual is true,
 	// then the Properties are associated with the backing sst.
@@ -2132,11 +2147,8 @@ func (d *DB) SSTables(opts ...SSTablesOption) ([][]SSTableInfo, error) {
 		fn(opt)
 	}
 
-	if opt.withApproximateSpanBytes && !opt.withProperties {
-		return nil, errors.Errorf("Cannot use WithApproximateSpanBytes without WithProperties option.")
-	}
 	if opt.withApproximateSpanBytes && (opt.start == nil || opt.end == nil) {
-		return nil, errors.Errorf("Cannot use WithApproximateSpanBytes without WithKeyRangeFilter option.")
+		return nil, errors.Errorf("cannot use WithApproximateSpanBytes without WithKeyRangeFilter option")
 	}
 
 	// Grab and reference the current readState.
@@ -2197,25 +2209,15 @@ func (d *DB) SSTables(opts ...SSTablesOption) ([][]SSTableInfo, error) {
 			}
 
 			if opt.withApproximateSpanBytes {
-				var spanBytes uint64
 				if m.ContainedWithinSpan(d.opts.Comparer.Compare, opt.start, opt.end) {
-					spanBytes = m.Size
+					destTables[j].ApproximateSpanBytes = m.Size
 				} else {
 					size, err := d.tableCache.estimateSize(m, opt.start, opt.end)
 					if err != nil {
 						return nil, err
 					}
-					spanBytes = size
-				}
-				propertiesCopy := *destTables[j].Properties
-
-				// Deep copy user properties so approximate span bytes can be added.
-				propertiesCopy.UserProperties = make(map[string]string, len(destTables[j].Properties.UserProperties)+1)
-				for k, v := range destTables[j].Properties.UserProperties {
-					propertiesCopy.UserProperties[k] = v
+					destTables[j].ApproximateSpanBytes = size
 				}
-				propertiesCopy.UserProperties["approximate-span-bytes"] = strconv.FormatUint(spanBytes, 10)
-				destTables[j].Properties = &propertiesCopy
 			}
 			j++
 		}


Test patch:
diff --git a/db_test.go b/db_test.go
index 5fd308b3..8a882ca4 100644
--- a/db_test.go
+++ b/db_test.go
@@ -1294,10 +1294,6 @@ func TestSSTablesWithApproximateSpanBytes(t *testing.T) {
 	require.NoError(t, d.Set([]byte("g"), nil, nil))
 	require.NoError(t, d.Flush())
 
-	// cannot use WithApproximateSpanBytes without WithProperties.
-	_, err = d.SSTables(WithKeyRangeFilter([]byte("a"), []byte("e")), WithApproximateSpanBytes())
-	require.Error(t, err)
-
 	// cannot use WithApproximateSpanBytes without WithKeyRangeFilter.
 	_, err = d.SSTables(WithProperties(), WithApproximateSpanBytes())
 	require.Error(t, err)
@@ -1307,13 +1303,11 @@ func TestSSTablesWithApproximateSpanBytes(t *testing.T) {
 
 	for _, levelTables := range tableInfos {
 		for _, table := range levelTables {
-			approximateSpanBytes, err := strconv.ParseInt(table.Properties.UserProperties["approximate-span-bytes"], 10, 64)
-			require.NoError(t, err)
 			if table.FileNum == 5 {
-				require.Equal(t, uint64(approximateSpanBytes), table.Size)
+				require.Equal(t, table.ApproximateSpanBytes, table.Size)
 			}
 			if table.FileNum == 7 {
-				require.Less(t, uint64(approximateSpanBytes), table.Size)
+				require.Less(t, table.ApproximateSpanBytes, table.Size)
 			}
 		}
 	}
diff --git a/internal/crdbtest/crdbtest.go b/internal/crdbtest/crdbtest.go
index 32d70195..df037310 100644
--- a/internal/crdbtest/crdbtest.go
+++ b/internal/crdbtest/crdbtest.go
@@ -154,9 +154,12 @@ func DecodeTimestamp(mvccKey []byte) ([]byte, []byte, uint64, uint32) {
 
 // Split implements base.Split for CockroachDB keys.
 func Split(key []byte) int {
+	// TODO(radu): Pebble sometimes passes empty "keys" and we have to tolerate
+	// them until we fix that.
 	if len(key) == 0 {
-		panic(errors.AssertionFailedf("empty key"))
+		return 0
 	}
+
 	// Last byte is the version length + 1 when there is a version, else it is
 	// 0.
 	versionLen := int(key[len(key)-1])
@@ -169,8 +172,10 @@ func Split(key []byte) int {
 // Compare compares cockroach keys, including the version (which could be MVCC
 // timestamps).
 func Compare(a, b []byte) int {
+	// TODO(radu): Pebble sometimes passes empty "keys" and we have to tolerate
+	// them until we fix that.
 	if len(a) == 0 || len(b) == 0 {
-		panic(errors.AssertionFailedf("malformed key: %x, %x", a, b))
+		return cmp.Compare(len(a), len(b))
 	}
 
 	// NB: For performance, this routine manually splits the key into the
@@ -219,11 +224,13 @@ func CompareSuffixes(a, b []byte) int {
 
 // Equal implements base.Equal for Cockroach keys.
 func Equal(a, b []byte) bool {
+	// TODO(radu): Pebble sometimes passes empty "keys" and we have to tolerate
+	// them until we fix that.
+	if len(a) == 0 || len(b) == 0 {
+		return len(a) == len(b)
+	}
 	aEnd := len(a) - 1
 	bEnd := len(b) - 1
-	if aEnd < 0 || bEnd < 0 {
-		panic("empty keys")
-	}
 
 	// Last byte is the version length + 1 when there is a version,
 	// else it is 0.