cli: Add 'pebble' debug command to run Pebble tool commands

pkg/cli/debug.go

@@ -53,6 +53,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/sysutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/cockroachdb/pebble"
+	"github.com/cockroachdb/pebble/tool"
 	"github.com/gogo/protobuf/jsonpb"
 	"github.com/kr/pretty"
 	"github.com/pkg/errors"
@@ -758,6 +760,14 @@ https://github.com/facebook/rocksdb/wiki/Administration-and-Data-Access-Tool#ldb
 	},
 }
 
+var debugPebbleCmd = &cobra.Command{
+	Use:   "pebble [command]",
+	Short: "run a Pebble introspection tool command",
+	Long: `
+Allows the use of pebble tools, such as to introspect manifests, SSTables, etc.
+`,
+}
+
 var debugSSTDumpCmd = &cobra.Command{
 	Use:   "sst_dump",
 	Short: "run the RocksDB 'sst_dump' tool",
@@ -1318,6 +1328,22 @@ process that has failed and cannot restart.
 func init() {
 	DebugCmd.AddCommand(debugCmds...)
 
+	pebbleTool := tool.New()
+	// To be able to read Cockroach-written RocksDB manifests/SSTables, comparator
+	// and merger functions must be specified to pebble that match the ones used
+	// to write those files.
+	//
+	// TODO(itsbilal): Port the Cockroach merger over from libroach/merge.cc to go
+	// and use that here. Until this happens, some data (eg. timeseries) will be
+	// printed incorrectly by this tool: it will be concatenated instead of being
+	// properly merged.
+	merger := *pebble.DefaultMerger
+	merger.Name = "cockroach_merge_operator"
+	pebbleTool.RegisterMerger(&merger)
+	pebbleTool.RegisterComparer(engine.MVCCComparer)
+	debugPebbleCmd.AddCommand(pebbleTool.Commands...)
+	DebugCmd.AddCommand(debugPebbleCmd)
+
 	f := debugSyncBenchCmd.Flags()
 	f.IntVarP(&syncBenchOpts.Concurrency, "concurrency", "c", syncBenchOpts.Concurrency,
 		"number of concurrent writers")

pkg/storage/bulk/sst_writer.go

@@ -30,39 +30,6 @@ type SSTWriter struct {
 	scratch  []byte
 }
 
-var mvccComparer = &pebble.Comparer{
-	Compare: engine.MVCCKeyCompare,
-	AbbreviatedKey: func(k []byte) uint64 {
-		key, _, ok := enginepb.SplitMVCCKey(k)
-		if !ok {
-			return 0
-		}
-		return pebble.DefaultComparer.AbbreviatedKey(key)
-	},
-
-	Separator: func(dst, a, b []byte) []byte {
-		return append(dst, a...)
-	},
-
-	Successor: func(dst, a []byte) []byte {
-		return append(dst, a...)
-	},
-	Split: func(k []byte) int {
-		if len(k) == 0 {
-			return len(k)
-		}
-		// This is similar to what enginepb.SplitMVCCKey does.
-		tsLen := int(k[len(k)-1])
-		keyPartEnd := len(k) - 1 - tsLen
-		if keyPartEnd < 0 {
-			return len(k)
-		}
-		return keyPartEnd
-	},
-
-	Name: "cockroach_comparator",
-}
-
 // timeboundPropCollector implements a property collector for MVCC Timestamps.
 // Its behavior matches TimeBoundTblPropCollector in table_props.cc.
 type timeboundPropCollector struct {
@@ -126,7 +93,7 @@ var pebbleOpts = func() *pebble.Options {
 	merger.Name = "nullptr"
 	opts := &pebble.Options{
 		TableFormat: pebble.TableFormatLevelDB,
-		Comparer:    mvccComparer,
+		Comparer:    engine.MVCCComparer,
 		Merger:      &merger,
 	}
 	opts.EnsureDefaults()

pkg/storage/engine/mvcc.go

@@ -26,7 +26,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
-	"github.com/pkg/errors"
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble"
 )
 
 const (
@@ -123,6 +124,11 @@ func (k MVCCKey) String() string {
 	return fmt.Sprintf("%s/%s", k.Key, k.Timestamp)
 }
 
+// Format implements the fmt.Formatter interface.
+func (k MVCCKey) Format(f fmt.State, c rune) {
+	fmt.Fprintf(f, "%s/%s", k.Key, k.Timestamp)
+}
+
 // Len returns the size of the MVCCKey when encoded. Implements the
 // pebble.Encodeable interface.
 //
@@ -1861,6 +1867,68 @@ func mvccInitPutUsingIter(
 		})
 }
 
+// mvccKeyFormatter is an fmt.Formatter for MVCC Keys.
+type mvccKeyFormatter struct {
+	key MVCCKey
+	err error
+}
+
+var _ fmt.Formatter = mvccKeyFormatter{}
+
+// Format implements the fmt.Formatter interface.
+func (m mvccKeyFormatter) Format(f fmt.State, c rune) {
+	if m.err != nil {
+		errors.FormatError(m.err, f, c)
+		return
+	}
+	m.key.Format(f, c)
+}
+
+// MVCCComparer is a pebble.Comparer object that implements MVCC-specific
+// comparator settings for use with Pebble.
+//
+// TODO(itsbilal): Move this to a new file pebble.go.
+var MVCCComparer = &pebble.Comparer{
+	Compare: MVCCKeyCompare,
+	AbbreviatedKey: func(k []byte) uint64 {
+		key, _, ok := enginepb.SplitMVCCKey(k)
+		if !ok {
+			return 0
+		}
+		return pebble.DefaultComparer.AbbreviatedKey(key)
+	},
+
+	Format: func(k []byte) fmt.Formatter {
+		decoded, err := DecodeMVCCKey(k)
+		if err != nil {
+			return mvccKeyFormatter{err: err}
+		}
+		return mvccKeyFormatter{key: decoded}
+	},
+
+	Separator: func(dst, a, b []byte) []byte {
+		return append(dst, a...)
+	},
+
+	Successor: func(dst, a []byte) []byte {
+		return append(dst, a...)
+	},
+	Split: func(k []byte) int {
+		if len(k) == 0 {
+			return len(k)
+		}
+		// This is similar to what enginepb.SplitMVCCKey does.
+		tsLen := int(k[len(k)-1])
+		keyPartEnd := len(k) - 1 - tsLen
+		if keyPartEnd < 0 {
+			return len(k)
+		}
+		return keyPartEnd
+	},
+
+	Name: "cockroach_comparator",
+}
+
 // MVCCMerge implements a merge operation. Merge adds integer values,
 // concatenates undifferentiated byte slice values, and efficiently
 // combines time series observations if the roachpb.Value tag value