From 84681dd5c1f2341bdcac204c36e959a8fb980155 Mon Sep 17 00:00:00 2001
From: Andrei Matei <andrei@cockroachlabs.com>
Date: Tue, 19 May 2020 17:26:24 -0400
Subject: [PATCH] cli: limit flag for debug range-data / debug keys

There was a teasing maxResults field already, completely unused.

Release note: None
---
 pkg/cli/cliflags/flags.go |  5 +++++
 pkg/cli/context.go        |  4 ++--
 pkg/cli/debug.go          | 15 ++++++++++++++-
 pkg/cli/flags.go          |  2 ++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/pkg/cli/cliflags/flags.go b/pkg/cli/cliflags/flags.go
index 56a378ea42c7..aa9879c3d1ae 100644
--- a/pkg/cli/cliflags/flags.go
+++ b/pkg/cli/cliflags/flags.go
@@ -885,6 +885,11 @@ Exclusive end key and format as [<format>:]<key>. Supported formats: raw, hex,
 human, rangeID. The raw format supports escaped text. For example, "raw:\x01k"
 is the prefix for range local keys. The hex format takes an encoded MVCCKey.`}
 
+	Limit = FlagInfo{
+		Name:        "limit",
+		Description: `Maximum number of keys to return.`,
+	}
+
 	Values = FlagInfo{
 		Name:        "values",
 		Description: `Print values along with their associated key.`,
diff --git a/pkg/cli/context.go b/pkg/cli/context.go
index 15c13eda1850..c115704f451a 100644
--- a/pkg/cli/context.go
+++ b/pkg/cli/context.go
@@ -105,7 +105,7 @@ func initCLIDefaults() {
 	debugCtx.replicated = false
 	debugCtx.inputFile = ""
 	debugCtx.printSystemConfig = false
-	debugCtx.maxResults = 1000
+	debugCtx.maxResults = 0
 	debugCtx.ballastSize = base.SizeSpec{InBytes: 1000000000}
 
 	serverCfg.GoroutineDumpDirName = ""
@@ -306,7 +306,7 @@ var debugCtx struct {
 	inputFile         string
 	ballastSize       base.SizeSpec
 	printSystemConfig bool
-	maxResults        int64
+	maxResults        int
 }
 
 // startCtx captures the command-line arguments for the `start` command.
diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go
index 3f1fd1216638..14a89906fa27 100644
--- a/pkg/cli/debug.go
+++ b/pkg/cli/debug.go
@@ -202,7 +202,15 @@ func runDebugKeys(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	return db.Iterate(debugCtx.startKey.Key, debugCtx.endKey.Key, printer)
+	results := 0
+	return db.Iterate(debugCtx.startKey.Key, debugCtx.endKey.Key, func(kv storage.MVCCKeyValue) (bool, error) {
+		done, err := printer(kv)
+		if done || err != nil {
+			return done, err
+		}
+		results++
+		return results == debugCtx.maxResults, nil
+	})
 }
 
 func runDebugBallast(cmd *cobra.Command, args []string) error {
@@ -292,6 +300,7 @@ func runDebugRangeData(cmd *cobra.Command, args []string) error {
 
 	iter := rditer.NewReplicaDataIterator(&desc, db, debugCtx.replicated, false /* seekEnd */)
 	defer iter.Close()
+	results := 0
 	for ; ; iter.Next() {
 		if ok, err := iter.Valid(); err != nil {
 			return err
@@ -302,6 +311,10 @@ func runDebugRangeData(cmd *cobra.Command, args []string) error {
 			Key:   iter.Key(),
 			Value: iter.Value(),
 		})
+		results++
+		if results == debugCtx.maxResults {
+			break
+		}
 	}
 	return nil
 }
diff --git a/pkg/cli/flags.go b/pkg/cli/flags.go
index 5c84c9a15fac..c72a77edb9da 100644
--- a/pkg/cli/flags.go
+++ b/pkg/cli/flags.go
@@ -686,12 +686,14 @@ func init() {
 		f := debugKeysCmd.Flags()
 		VarFlag(f, (*mvccKey)(&debugCtx.startKey), cliflags.From)
 		VarFlag(f, (*mvccKey)(&debugCtx.endKey), cliflags.To)
+		IntFlag(f, &debugCtx.maxResults, cliflags.Limit, debugCtx.maxResults)
 		BoolFlag(f, &debugCtx.values, cliflags.Values, debugCtx.values)
 		BoolFlag(f, &debugCtx.sizes, cliflags.Sizes, debugCtx.sizes)
 	}
 	{
 		f := debugRangeDataCmd.Flags()
 		BoolFlag(f, &debugCtx.replicated, cliflags.Replicated, debugCtx.replicated)
+		IntFlag(f, &debugCtx.maxResults, cliflags.Limit, debugCtx.maxResults)
 	}
 	{
 		f := debugGossipValuesCmd.Flags()