mvcc: reduce count-only range overhead

CHANGELOG-3.5.md

@@ -101,6 +101,8 @@ Note that any `etcd_debugging_*` metrics are experimental and subject to change.
 - Improve [compaction performance when latest index is greater than 1-million](https://github.com/etcd-io/etcd/pull/11734).
 - [Refactor consistentindex](https://github.com/etcd-io/etcd/pull/11699).
 - [Add log when etcdserver failed to apply command](https://github.com/etcd-io/etcd/pull/11670).
+- Improve [count-only range performance](https://github.com/etcd-io/etcd/pull/11771).
+
 
 ### Package `embed`
 

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/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