Merge pull request #15099 from fuweid/backport-11771-11743-pr-to-3.4

[3.4] mvcc: reduce count-only range overhead

etcdctl/ctlv3/command/get_command.go

@@ -31,6 +31,7 @@ var (
 	getFromKey     bool
 	getRev         int64
 	getKeysOnly    bool
+	getCountOnly   bool
 	printValueOnly bool
 )
 
@@ -50,6 +51,7 @@ func NewGetCommand() *cobra.Command {
 	cmd.Flags().BoolVar(&getFromKey, "from-key", false, "Get keys that are greater than or equal to the given key using byte compare")
 	cmd.Flags().Int64Var(&getRev, "rev", 0, "Specify the kv revision")
 	cmd.Flags().BoolVar(&getKeysOnly, "keys-only", false, "Get only the keys")
+	cmd.Flags().BoolVar(&getCountOnly, "count-only", false, "Get only the count")
 	cmd.Flags().BoolVar(&printValueOnly, "print-value-only", false, `Only write values when using the "simple" output format`)
 	return cmd
 }
@@ -64,6 +66,12 @@ func getCommandFunc(cmd *cobra.Command, args []string) {
 		ExitWithError(ExitError, err)
 	}
 
+	if getCountOnly {
+		if _, fields := display.(*fieldsPrinter); !fields {
+			ExitWithError(ExitBadArgs, fmt.Errorf("--count-only is only for `--write-out=fields`"))
+		}
+	}
+
 	if printValueOnly {
 		dp, simple := (display).(*simplePrinter)
 		if !simple {
@@ -83,6 +91,10 @@ func getGetOp(args []string) (string, []clientv3.OpOption) {
 		ExitWithError(ExitBadArgs, fmt.Errorf("`--prefix` and `--from-key` cannot be set at the same time, choose one"))
 	}
 
+	if getKeysOnly && getCountOnly {
+		ExitWithError(ExitBadArgs, fmt.Errorf("`--keys-only` and `--count-only` cannot be set at the same time, choose one"))
+	}
+
 	opts := []clientv3.OpOption{}
 	switch getConsistency {
 	case "s":
@@ -159,5 +171,9 @@ func getGetOp(args []string) (string, []clientv3.OpOption) {
 		opts = append(opts, clientv3.WithKeysOnly())
 	}
 
+	if getCountOnly {
+		opts = append(opts, clientv3.WithCountOnly())
+	}
+
 	return key, opts
 }

mvcc/index.go

@@ -26,6 +26,7 @@ type index interface {
 	Get(key []byte, atRev int64) (rev, created revision, ver int64, err error)
 	Range(key, end []byte, atRev int64) ([][]byte, []revision)
 	Revisions(key, end []byte, atRev int64) []revision
+	CountRevisions(key, end []byte, atRev int64) int
 	Put(key []byte, rev revision)
 	Tombstone(key []byte, rev revision) error
 	RangeSince(key, end []byte, rev int64) []revision
@@ -119,6 +120,23 @@ func (ti *treeIndex) Revisions(key, end []byte, atRev int64) (revs []revision) {
 	return revs
 }
 
+func (ti *treeIndex) CountRevisions(key, end []byte, atRev int64) int {
+	if end == nil {
+		_, _, _, err := ti.Get(key, atRev)
+		if err != nil {
+			return 0
+		}
+		return 1
+	}
+	total := 0
+	ti.visit(key, end, func(ki *keyIndex) {
+		if _, _, _, err := ki.get(ti.lg, atRev); err == nil {
+			total++
+		}
+	})
+	return total
+}
+
 func (ti *treeIndex) Range(key, end []byte, atRev int64) (keys [][]byte, revs []revision) {
 	if end == nil {
 		rev, _, _, err := ti.Get(key, atRev)

mvcc/index_test.go

@@ -194,6 +194,80 @@ func TestIndexRangeSince(t *testing.T) {
 	}
 }
 
+func TestIndexRevision(t *testing.T) {
+	allKeys := [][]byte{[]byte("foo"), []byte("foo1"), []byte("foo2"), []byte("foo2"), []byte("foo1"), []byte("foo")}
+	allRevs := []revision{{main: 1}, {main: 2}, {main: 3}, {main: 4}, {main: 5}, {main: 6}}
+
+	ti := newTreeIndex(zap.NewExample())
+	for i := range allKeys {
+		ti.Put(allKeys[i], allRevs[i])
+	}
+
+	tests := []struct {
+		key, end []byte
+		atRev    int64
+		wrevs    []revision
+		wcounts  int
+	}{
+		// single key that not found
+		{
+			[]byte("bar"), nil, 6, nil, 0,
+		},
+		// single key that found
+		{
+			[]byte("foo"), nil, 6, []revision{{main: 6}}, 1,
+		},
+		// various range keys, fixed atRev
+		{
+			[]byte("foo"), []byte("foo1"), 6, []revision{{main: 6}}, 1,
+		},
+		{
+			[]byte("foo"), []byte("foo2"), 6, []revision{{main: 6}, {main: 5}}, 2,
+		},
+		{
+			[]byte("foo"), []byte("fop"), 6, []revision{{main: 6}, {main: 5}, {main: 4}}, 3,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2,
+		},
+		{
+			[]byte("foo2"), []byte("fop"), 6, []revision{{main: 4}}, 1,
+		},
+		{
+			[]byte("foo3"), []byte("fop"), 6, nil, 0,
+		},
+		// fixed range keys, various atRev
+		{
+			[]byte("foo1"), []byte("fop"), 1, nil, 0,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 2, []revision{{main: 2}}, 1,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 3, []revision{{main: 2}, {main: 3}}, 2,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 4, []revision{{main: 2}, {main: 4}}, 2,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 5, []revision{{main: 5}, {main: 4}}, 2,
+		},
+		{
+			[]byte("foo1"), []byte("fop"), 6, []revision{{main: 5}, {main: 4}}, 2,
+		},
+	}
+	for i, tt := range tests {
+		revs := ti.Revisions(tt.key, tt.end, tt.atRev)
+		if !reflect.DeepEqual(revs, tt.wrevs) {
+			t.Errorf("#%d: revs = %+v, want %+v", i, revs, tt.wrevs)
+		}
+		count := ti.CountRevisions(tt.key, tt.end, tt.atRev)
+		if count != tt.wcounts {
+			t.Errorf("#%d: count = %d, want %v", i, count, tt.wcounts)
+		}
+	}
+}
+
 func TestIndexCompactAndKeep(t *testing.T) {
 	maxRev := int64(20)
 	tests := []struct {

mvcc/kvstore_test.go

@@ -941,6 +941,11 @@ func (i *fakeIndex) Revisions(key, end []byte, atRev int64) []revision {
 	return rev
 }
 
+func (i *fakeIndex) CountRevisions(key, end []byte, atRev int64) int {
+	_, rev := i.Range(key, end, atRev)
+	return len(rev)
+}
+
 func (i *fakeIndex) Get(key []byte, atRev int64) (rev, created revision, ver int64, err error) {
 	i.Recorder.Record(testutil.Action{Name: "get", Params: []interface{}{key, atRev}})
 	r := <-i.indexGetRespc

mvcc/kvstore_txn.go

@@ -125,15 +125,16 @@ func (tr *storeTxnRead) rangeKeys(key, end []byte, curRev int64, ro RangeOptions
 	if rev < tr.s.compactMainRev {
 		return &RangeResult{KVs: nil, Count: -1, Rev: 0}, ErrCompacted
 	}
-
+	if ro.Count {
+		total := tr.s.kvindex.CountRevisions(key, end, rev)
+		tr.trace.Step("count revisions from in-memory index tree")
+		return &RangeResult{KVs: nil, Count: total, Rev: curRev}, nil
+	}
 	revpairs := tr.s.kvindex.Revisions(key, end, rev)
 	tr.trace.Step("range keys from in-memory index tree")
 	if len(revpairs) == 0 {
 		return &RangeResult{KVs: nil, Count: 0, Rev: curRev}, nil
 	}
-	if ro.Count {
-		return &RangeResult{KVs: nil, Count: len(revpairs), Rev: curRev}, nil
-	}
 
 	limit := int(ro.Limit)
 	if limit <= 0 || limit > len(revpairs) {

tests/e2e/ctl_v3_kv_test.go

@@ -41,9 +41,10 @@ func TestCtlV3GetPeerTLS(t *testing.T)       { testCtl(t, getTest, withCfg(confi
 func TestCtlV3GetTimeout(t *testing.T)       { testCtl(t, getTest, withDialTimeout(0)) }
 func TestCtlV3GetQuorum(t *testing.T)        { testCtl(t, getTest, withQuorum()) }
 
-func TestCtlV3GetFormat(t *testing.T)   { testCtl(t, getFormatTest) }
-func TestCtlV3GetRev(t *testing.T)      { testCtl(t, getRevTest) }
-func TestCtlV3GetKeysOnly(t *testing.T) { testCtl(t, getKeysOnlyTest) }
+func TestCtlV3GetFormat(t *testing.T)    { testCtl(t, getFormatTest) }
+func TestCtlV3GetRev(t *testing.T)       { testCtl(t, getRevTest) }
+func TestCtlV3GetKeysOnly(t *testing.T)  { testCtl(t, getKeysOnlyTest) }
+func TestCtlV3GetCountOnly(t *testing.T) { testCtl(t, getCountOnlyTest) }
 
 func TestCtlV3Del(t *testing.T)          { testCtl(t, delTest) }
 func TestCtlV3DelNoTLS(t *testing.T)     { testCtl(t, delTest, withCfg(configNoTLS)) }
@@ -235,6 +236,44 @@ func getKeysOnlyTest(cx ctlCtx) {
 	}
 }
 
+func getCountOnlyTest(cx ctlCtx) {
+	cmdArgs := append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...)
+	if err := spawnWithExpects(cmdArgs, "\"Count\" : 0"); err != nil {
+		cx.t.Fatal(err)
+	}
+	if err := ctlV3Put(cx, "key", "val", ""); err != nil {
+		cx.t.Fatal(err)
+	}
+	cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...)
+	if err := spawnWithExpects(cmdArgs, "\"Count\" : 1"); err != nil {
+		cx.t.Fatal(err)
+	}
+	if err := ctlV3Put(cx, "key1", "val", ""); err != nil {
+		cx.t.Fatal(err)
+	}
+	if err := ctlV3Put(cx, "key1", "val", ""); err != nil {
+		cx.t.Fatal(err)
+	}
+	cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...)
+	if err := spawnWithExpects(cmdArgs, "\"Count\" : 2"); err != nil {
+		cx.t.Fatal(err)
+	}
+	if err := ctlV3Put(cx, "key2", "val", ""); err != nil {
+		cx.t.Fatal(err)
+	}
+	cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key", "--prefix", "--write-out=fields"}...)
+	if err := spawnWithExpects(cmdArgs, "\"Count\" : 3"); err != nil {
+		cx.t.Fatal(err)
+	}
+	expected := []string{
+		"\"Count\" : 3",
+	}
+	cmdArgs = append(cx.PrefixArgs(), []string{"get", "--count-only", "key3", "--prefix", "--write-out=fields"}...)
+	if err := spawnWithExpects(cmdArgs, expected...); err == nil {
+		cx.t.Fatal(err)
+	}
+}
+
 func delTest(cx ctlCtx) {
 	tests := []struct {
 		puts []kv